Add native FP8 expert support to REAP - #3003
Conversation
Signed-off-by: Ilya Boytsov <ieboytsov@nebius.com>
|
Important Review skippedAuto reviews are disabled on this repository. 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:
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 |
Merge Protections🔴 2 of 2 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.
🔴 Require two reviewsWaiting for
This rule is failing.PRs labelled "two-reviews" must have at least two approving reviews before merging.
|
![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 introduces REAP (Representation-based Expert Pruning) support for native fine-grained FP8 experts in Transformers. It adds the FP8PrunableExperts class, defines the ReapPrunableExpertsProtocol interface, and updates the REAPPruningModifier and utility functions to handle both standard 2D linear experts and packed FP8 expert containers. Comprehensive unit tests are also added to verify the new pruning lifecycle. The review feedback highlights three key areas for improvement: wrapping the FP8Experts import in a try-except block to maintain backward compatibility with older transformers versions, adding defensive checks in _prune_expert_tensor to handle None values safely, and converting the expert_indices tensor to a Python list to prevent host-device synchronization bottlenecks during iteration.
| import torch | ||
| import torch.nn as nn | ||
| from compressed_tensors import align_module_device | ||
| from transformers.integrations.finegrained_fp8 import FP8Experts |
There was a problem hiding this comment.
Unconditionally importing FP8Experts from transformers.integrations.finegrained_fp8 will cause a ModuleNotFoundError or ImportError on environments running older versions of transformers that do not support fine-grained FP8. Since this module is imported unconditionally in utils.py, it breaks backward compatibility for all users of the library. Wrapping this import in a try-except block and defining a fallback dummy class with a dummy forward method ensures backward compatibility.
| from transformers.integrations.finegrained_fp8 import FP8Experts | |
| try: | |
| from transformers.integrations.finegrained_fp8 import FP8Experts | |
| except ImportError: | |
| class FP8Experts: | |
| def forward(self, *args, **kwargs): | |
| pass |
| retained_t = torch.tensor(retained, dtype=torch.long) | ||
| with align_module_device(holder): | ||
| value = getattr(holder, name) | ||
| retained_t = retained_t.to(value.device) | ||
| pruned = value.detach()[retained_t].contiguous() | ||
| requires_grad = value.requires_grad |
There was a problem hiding this comment.
In _prune_expert_tensor, if a parameter or buffer exists but its value is None (which is common for optional parameters/buffers like bias or e_score_correction_bias in some configurations), getattr(holder, name) will return None. Attempting to access value.device or value.requires_grad will then raise an AttributeError. Adding a defensive check to return early if value is None prevents this potential runtime crash.
value = getattr(holder, name, None)
if value is None:
return
retained_t = torch.tensor(retained, dtype=torch.long)
with align_module_device(holder):
retained_t = retained_t.to(value.device)
pruned = value.detach()[retained_t].contiguous()
requires_grad = getattr(value, "requires_grad", False)| if calibrate_all: | ||
| expert_indices = range(self.num_experts) | ||
| else: | ||
| expert_indices = ( | ||
| torch.greater(expert_mask.sum(dim=(-1, -2)), 0) | ||
| .nonzero(as_tuple=False) | ||
| .view(-1) | ||
| ) |
There was a problem hiding this comment.
Iterating directly over expert_indices when it is a PyTorch GPU tensor (returned by .nonzero().view(-1)) causes host-device synchronization on every single iteration of the loop. This introduces a significant performance bottleneck during calibration/inference. Converting the tensor to a standard Python list of integers using .tolist() before iterating avoids these synchronizations and dramatically improves execution speed.
| if calibrate_all: | |
| expert_indices = range(self.num_experts) | |
| else: | |
| expert_indices = ( | |
| torch.greater(expert_mask.sum(dim=(-1, -2)), 0) | |
| .nonzero(as_tuple=False) | |
| .view(-1) | |
| ) | |
| if calibrate_all: | |
| expert_indices = list(range(self.num_experts)) | |
| else: | |
| expert_indices = ( | |
| torch.greater(expert_mask.sum(dim=(-1, -2)), 0) | |
| .nonzero(as_tuple=False) | |
| .view(-1) | |
| .tolist() | |
| ) |
|
👋 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. |
1 participant
Summary:
Currently REAP can't work with fp8 checkpoints because
LinearExperts2Dassumes a checkpoint packs experts as a list ofnn.linearinstances which is true for a lot of models but experts in FP8 checkpoints (at least some of what I saw) are packed all together into a fused containerFP8ExpertssoLinearExperts2Dcan't iterate them. This PR fixes the problem.The only way to use fp8 checkpoints with REAP now is to dequantize chpt in the beginning and run calibration with bf/fp16 weights. This also saved pruned model in bf/fp16 rather than in original fp8 which is ambigious.
Adds REAP support for packed native FP8 MoE experts, including activation-norm collection and structural pruning while preserving FP8 weights and scales. Updates router and correction-bias pruning and adds coverage for native FP8 expert detection, calibration, and pruning.
Test plan
Validated with minimax-m2.5 model and 38 REAP tests plus Ruff checks