[WIP] [MoE] No-copy 3D loading - #2941
Conversation
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
|
Important Review skippedIgnore keyword(s) in the title. ⛔ Ignored keywords (3)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review. Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the MoE linearization and conversion mapping logic, introducing dynamic mapping registration, support for the inkling_mm_model architecture, and a zero-copy weight assignment mechanism in from_experts_module to optimize offloading. However, several critical issues need to be addressed: the use of the unsupported log_once parameter in loguru warnings will cause runtime crashes; registering save_mappings instead of mappings during saving discards other weight conversions; and the zero-copy logic fails to handle biases and transposed experts. Additionally, please restore the commented-out test cases and the load_quantizable_moe context manager to maintain test coverage.
| logger.warning( | ||
| "Linearized model performs a weight conversion during loading. This " | ||
| f"may lead to longer load times\n{converter}" | ||
| "may lead to longer load times", | ||
| log_once=True, | ||
| ) |
There was a problem hiding this comment.
The log_once parameter is not supported by loguru.logger.warning and will raise a TypeError at runtime. Please remove log_once=True to prevent crashes.
logger.warning(
"Linearized model performs a weight conversion during loading. This "
"may lead to longer load times"
)| logger.warning( | ||
| "Linearized model performs a weight conversion during saving. This " | ||
| f"may lead to longer save times\n{converter}" | ||
| "may lead to longer save times", | ||
| log_once=True, | ||
| ) |
There was a problem hiding this comment.
The log_once parameter is not supported by loguru.logger.warning and will raise a TypeError at runtime. Please remove log_once=True to prevent crashes.
logger.warning(
"Linearized model performs a weight conversion during saving. This "
"may lead to longer save times"
)| model._weight_conversions = save_mappings | ||
| register_checkpoint_conversion_mapping(model_type, save_mappings, overwrite=True) |
There was a problem hiding this comment.
Assigning and registering save_mappings instead of mappings will cause the model to lose all other weight conversions during saving. Please use mappings instead.
| model._weight_conversions = save_mappings | |
| register_checkpoint_conversion_mapping(model_type, save_mappings, overwrite=True) | |
| model._weight_conversions = mappings | |
| register_checkpoint_conversion_mapping(model_type, mappings, overwrite=True) |
| # zero copy: remove weight to avoid offloading | ||
| for index in range(self.num_experts): | ||
| expert: ExpertMLP = self[index] | ||
| expert.copy_from_experts_module(experts, index) | ||
| if hasattr(expert, "gate_proj"): | ||
| del expert.gate_proj.weight | ||
| del expert.up_proj.weight | ||
| del expert.down_proj.weight | ||
|
|
||
| # copy offloading from original | ||
| offload_kwargs = get_cache_init_kwargs(experts) | ||
| for module in self.modules(): | ||
| offload_module(module, **offload_kwargs) | ||
|
|
||
| # zero copy: assign linear weights as a view | ||
| with disable_onloading(): | ||
| for index in range(self.num_experts): | ||
| expert: ExpertMLP = self[index] | ||
| if hasattr(expert, "gate_proj"): | ||
| expert.gate_proj._parameters["weight"] = experts.gate_up_proj[ | ||
| index, : self.intermediate_size | ||
| ] | ||
| expert.up_proj._parameters["weight"] = experts.gate_up_proj[ | ||
| index, self.intermediate_size : | ||
| ] | ||
| expert.down_proj._parameters["weight"] = experts.down_proj[index] |
There was a problem hiding this comment.
The new from_experts_module implementation has two major issues:
- It completely ignores biases. If the model has biases enabled, they will remain on the "meta" device or uninitialized, causing runtime crashes or incorrect outputs.
- It ignores
experts.is_transposed. If the experts are transposed, slicing and assigning the weights directly will result in incorrect shapes and weights.
We should update this block to correctly handle both transposition and biases.
# zero copy: remove weight and bias to avoid offloading
for index in range(self.num_experts):
expert: ExpertMLP = self[index]
if hasattr(expert, "gate_proj"):
del expert.gate_proj.weight
if experts.has_bias:
del expert.gate_proj.bias
del expert.up_proj.weight
if experts.has_bias:
del expert.up_proj.bias
del expert.down_proj.weight
if experts.has_bias:
del expert.down_proj.bias
# copy offloading from original
offload_kwargs = get_cache_init_kwargs(experts)
for module in self.modules():
offload_module(module, **offload_kwargs)
# zero copy: assign linear weights and biases as a view
with disable_onloading():
for index in range(self.num_experts):
expert: ExpertMLP = self[index]
if not experts.is_transposed:
gate_up_proj = experts.gate_up_proj[index]
down_proj = experts.down_proj[index]
else:
gate_up_proj = experts.gate_up_proj[index].T
down_proj = experts.down_proj[index].T
if hasattr(expert, "gate_proj"):
expert.gate_proj._parameters["weight"] = gate_up_proj[: self.intermediate_size]
if experts.has_bias:
expert.gate_proj._parameters["bias"] = experts.gate_up_proj_bias[
index, : self.intermediate_size
]
expert.up_proj._parameters["weight"] = gate_up_proj[self.intermediate_size :]
if experts.has_bias:
expert.up_proj._parameters["bias"] = experts.gate_up_proj_bias[
index, self.intermediate_size :
]
expert.down_proj._parameters["weight"] = down_proj
if experts.has_bias:
expert.down_proj._parameters["bias"] = experts.down_proj_bias[index]There was a problem hiding this comment.
- I agree with the biases part
is_transposedis handled by the weight converter and at runtime. At this point, it's standardized
| # with load_quantizable_moe(model_cls): | ||
| from compressed_tensors.offload import set_onload_device | ||
|
|
||
| from llmcompressor.utils import load_context | ||
|
|
||
| with load_context(model_cls): | ||
| model2 = model_cls.from_pretrained( | ||
| model_stub, | ||
| device_map="auto_offload", | ||
| max_memory={}, | ||
| offload_folder="offload_folder", | ||
| ) | ||
| set_onload_device(model2, "cuda") | ||
| # with load_quantizable_moe(model_cls): | ||
| # model2 = model_cls.from_pretrained(model_stub, device_map="cuda") |
There was a problem hiding this comment.
The load_quantizable_moe context manager is commented out and replaced with load_context. This means the test is no longer verifying the behavior of load_quantizable_moe. Please restore the original context manager to ensure the feature is properly tested.
| # with load_quantizable_moe(model_cls): | |
| from compressed_tensors.offload import set_onload_device | |
| from llmcompressor.utils import load_context | |
| with load_context(model_cls): | |
| model2 = model_cls.from_pretrained( | |
| model_stub, | |
| device_map="auto_offload", | |
| max_memory={}, | |
| offload_folder="offload_folder", | |
| ) | |
| set_onload_device(model2, "cuda") | |
| # with load_quantizable_moe(model_cls): | |
| # model2 = model_cls.from_pretrained(model_stub, device_map="cuda") | |
| with load_quantizable_moe(model_cls): | |
| model2 = model_cls.from_pretrained(model_stub, device_map="cuda") |
| # ( | ||
| # "inference-optimization/DSV4-tiny-empty", | ||
| # [ | ||
| # "model.layers.0.mlp.experts.2.up_proj.weight", | ||
| # "model.layers.1.mlp.experts.0.gate_proj.weight", | ||
| # "model.layers.2.mlp.experts.1.down_proj.weight", | ||
| # ], | ||
| # ), | ||
| # ( | ||
| # "inference-optimization/Qwen3-1.6B-A0.9B", | ||
| # [ | ||
| # "model.layers.0.mlp.experts.2.up_proj.weight", | ||
| # "model.layers.1.mlp.experts.0.gate_proj.weight", | ||
| # "model.layers.2.mlp.experts.1.down_proj.weight", | ||
| # ], | ||
| # ), | ||
| # ( | ||
| # "inference-optimization/GLM-5.2-0.8B-A0.8B", | ||
| # [ | ||
| # "model.layers.2.mlp.experts.2.up_proj.weight", | ||
| # "model.layers.3.mlp.experts.0.gate_proj.weight", | ||
| # "model.layers.4.mlp.experts.1.down_proj.weight", | ||
| # ], | ||
| # ), |
There was a problem hiding this comment.
These test cases for DeepSeek V4, Qwen3, and GLM models have been commented out, which disables test coverage for these architectures. Please restore them to ensure no regressions are introduced.
(
"inference-optimization/DSV4-tiny-empty",
[
"model.layers.0.mlp.experts.2.up_proj.weight",
"model.layers.1.mlp.experts.0.gate_proj.weight",
"model.layers.2.mlp.experts.1.down_proj.weight",
],
),
(
"inference-optimization/Qwen3-1.6B-A0.9B",
[
"model.layers.0.mlp.experts.2.up_proj.weight",
"model.layers.1.mlp.experts.0.gate_proj.weight",
"model.layers.2.mlp.experts.1.down_proj.weight",
],
),
(
"inference-optimization/GLM-5.2-0.8B-A0.8B",
[
"model.layers.2.mlp.experts.2.up_proj.weight",
"model.layers.3.mlp.experts.0.gate_proj.weight",
"model.layers.4.mlp.experts.1.down_proj.weight",
],
),Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
|
The quality checks have failed. Please run |
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
|
The quality checks have failed. Please run |
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
|
The quality checks have failed. Please run |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require one maintainer reviewWaiting for any of
This rule is failing.All PRs must have at least one approving review from a maintainer before merging.
|
|
This pull request has merge conflicts that must be resolved before it can be |
|
Rather than add lots of complex code, let's just linearize moes on the fly. Users will have to eat the cost of up-front linearization if the sequential targets are smaller, but that should be an uncommon case |
2 participants
Purpose
TODO
Prerequisites
Changes
Testing
Before:
After:
I've confirmed that remaining runtime is purely cpu overhead