| |
| |
| |
| |
| @@ -27,7 +27,10 @@ import logging |
| import torch |
| from datasets import Dataset, DatasetDict, load_dataset |
| from transformers import AutoTokenizer, BatchEncoding |
| -from vllm import TokensPrompt |
| +try: |
| + from vllm import TokensPrompt |
| +except ImportError: # apex: only needed with return_vllm_tokens_prompt |
| + TokensPrompt = None |
| |
| |
| logger = logging.getLogger(__name__) |
| @@ -214,6 +217,8 @@ def _load_raw_dataset(dataset_name, split, subset=None): |
| return load_dataset( |
| "json", data_files={"train": file_url}, split="train", streaming=False |
| ) |
| + elif dataset_name.endswith((".json", ".jsonl")): # apex: a local dataset file |
| + return load_dataset("json", data_files=dataset_name, split="train") |
| else: |
| load_kwargs = {} |
| if subset is not None: |
| @@ -241,6 +246,8 @@ def load_category_batches( |
| |
| # load dataset processor |
| proc_cls = DATASET_REGISTRY.get(dataset_name) |
| + if proc_cls is None and dataset_name.endswith((".json", ".jsonl")): |
| + proc_cls = ApexCalibrationDataset |
| if proc_cls is None: |
| raise ValueError( |
| f"No DatasetProcessor registered for '{dataset_name}'. " |
| @@ -962,6 +969,24 @@ class SWESmithTrajectoriesDataset(ChatDatasetProcessor): |
| } |
| |
| |
| +class ApexCalibrationDataset(BaseDatasetProcessor): |
| + """apex-moe's calibration prompts with their datasets' answers (one JSON line per prompt): |
| + "messages" (user prompt + the dataset's answer) go through the chat template as in |
| + ChatDatasetProcessor; rows with "text" only (FineWeb-Edu passages) are tokenized raw as in |
| + LMDatasetProcessor.""" |
| + |
| + category_field: str = "source" |
| + |
| + @staticmethod |
| + def _map_fn(sample: dict[str, any]) -> dict[str, any]: |
| + return sample |
| + |
| + def _encode_sample(self, sample: dict) -> torch.Tensor: |
| + if sample.get(self.messages_field): |
| + return ChatDatasetProcessor._encode_sample(self, sample) |
| + return LMDatasetProcessor._encode_sample(self, sample) |
| + |
| + |
| DATASET_REGISTRY: dict[str, BaseDatasetProcessor] = { |
| "m-a-p/CodeFeedback-Filtered-Instruction": CodeFeedbackChatDataset, |
| "allenai/tulu-3-sft-mixture": TuluSFTMixtureChatDataset, |
| |
| |
| |
| |
| @@ -52,7 +52,6 @@ from reap.model_util import ( |
| patched_model_map, |
| get_super_expert_indices, |
| ) |
| -from reap.eval import run_evaluate |
| from reap.cluster_plots import plot_cluster_analysis |
| from reap.metrics import get_distance_fn |
| |
| @@ -118,8 +117,13 @@ def create_results_directory(model_name: str, dataset_name: str) -> pathlib.Path |
| def _setup_observer(model, obs_args): |
| """Create and return an MoETransformerObserver for the given model.""" |
| try: |
| + norm_topk_prob = getattr(model.config, "norm_topk_prob", None) |
| + if norm_topk_prob is None and model.__class__.__name__ == "Qwen3_5MoeForCausalLM": |
| + # apex: Qwen3.5/3.6's router always renormalizes its top-k weights (transformers |
| + # 5 Qwen3_5MoeTopKRouter) and its config has no norm_topk_prob |
| + norm_topk_prob = True |
| renormalize_router_weights = ( |
| - getattr(model.config, "norm_topk_prob", False) |
| + bool(norm_topk_prob) |
| and obs_args.renormalize_router_weights |
| ) |
| if renormalize_router_weights: |
| @@ -547,10 +551,11 @@ def smoke_test(model: nn.Module, tokenizer: AutoTokenizer): |
| return_tensors="pt", |
| add_generation_prompt=True, |
| tokenize=True, |
| + return_dict=True, # apex: transformers 5 returns a BatchEncoding either way |
| # enable_thinking=False, |
| ).to(model.device) |
| outputs = model.generate( |
| - inputs, |
| + **inputs, |
| max_new_tokens=50, |
| do_sample=True, |
| ) |
| @@ -829,6 +834,7 @@ def main(): |
| torch.cuda.empty_cache() |
| gc.collect() |
| model_args.model_name = merged_model_dir |
| + from reap.eval import run_evaluate # apex: lazy, pulls in vllm / lm_eval / evalplus |
| run_evaluate(model_args, merged_model_dir / "eval", eval_args, reap_args.seed) |
| |
| |
| |
| |
| |
| |
| @@ -104,6 +104,30 @@ MODEL_ATTRS = { |
| "num_experts": "num_experts", |
| "num_experts_per_tok": "num_experts_per_tok", |
| }, |
| + # apex: transformers 5 Qwen3.5/3.6 and Qwen4Exp (Qwen3.8): 3D experts, router `gate` |
| + "Qwen3_5MoeForCausalLM": { |
| + "moe_block": "mlp", |
| + "gate_proj": "gate_up_proj", |
| + "up_proj": "gate_up_proj", |
| + "down_proj": "down_proj", |
| + "experts": "experts", |
| + "fused": True, |
| + "router": "gate", |
| + "num_experts": "num_experts", |
| + "num_experts_per_tok": "num_experts_per_tok", |
| + }, |
| + # apex: transformers 5 Qwen3.5/3.6 and Qwen4Exp (Qwen3.8): 3D experts, router `gate` |
| + "Qwen4ExpForCausalLM": { |
| + "moe_block": "mlp", |
| + "gate_proj": "gate_up_proj", |
| + "up_proj": "gate_up_proj", |
| + "down_proj": "down_proj", |
| + "experts": "experts", |
| + "fused": True, |
| + "router": "gate", |
| + "num_experts": "num_experts", |
| + "num_experts_per_tok": "num_experts_per_tok", |
| + }, |
| "Glm4MoeForCausalLM": { |
| "moe_block": "mlp", |
| "gate_proj": "gate_proj", |
| |
| |
| |
| |
| @@ -220,6 +220,10 @@ class MoETransformerObserverConfig(BaseTransformerObserverHookConfig): |
| num_experts_attr_name: str = "num_experts" |
| top_k_attr_name: str = "top_k" |
| fused_experts: bool = False |
| + # apex: transformers 5 MoE blocks (Qwen3.5/3.6, Qwen4Exp): the block returns only its hidden |
| + # states, the router is `gate` (logits, top-k weights, top-k indices) and the experts are one |
| + # module with 3D `gate_up_proj` [E, 2I, H] and `down_proj` [E, H, I] |
| + fused_3d_experts: bool = False |
| distance_measure: str = "angular" |
| renormalize_router_weights: bool = False |
| record_pruning_metrics_only: bool = False |
| @@ -328,7 +332,7 @@ class MoETransformerObserver(BaseTransformerObserver): |
| |
| @torch.no_grad() |
| def _hook_fn(module, args, output): |
| - if not len(output) >= 2: |
| + if not self.hook_config.fused_3d_experts and not len(output) >= 2: |
| raise ValueError( |
| f"Expected output of module {module.__class__.__name__} at layer " |
| f"{layer_number} to be a tuple of at least length 2, got {len(output)}." |
| @@ -350,7 +354,33 @@ class MoETransformerObserver(BaseTransformerObserver): |
| |
| activations = torch.zeros((num_experts, *flat_input.shape), device=device) |
| |
| - if self.hook_config.fused_experts: |
| + if self.hook_config.fused_3d_experts: |
| + # as the loop below: router logits, top-k by logit, every expert's unweighted |
| + # output on every token |
| + router_logits = torch.nn.functional.linear(flat_input, module.gate.weight) |
| + _, selected_experts = torch.topk(router_logits, top_k, dim=-1) |
| + experts = module.experts |
| + fp8 = hasattr(experts, "gate_up_proj_scale_inv") # transformers FP8Experts |
| + for idx in range(num_experts): |
| + if fp8: # the module's own block-FP8 GEMMs, as its forward computes them |
| + out = experts.linear( |
| + flat_input, experts.gate_up_proj[idx], experts.gate_up_proj_scale_inv[idx] |
| + ) |
| + out = experts.linear( |
| + experts._apply_gate(out), |
| + experts.down_proj[idx], |
| + experts.down_proj_scale_inv[idx], |
| + ) |
| + else: |
| + gate, up = torch.nn.functional.linear( |
| + flat_input, experts.gate_up_proj[idx] |
| + ).chunk(2, dim=-1) |
| + out = torch.nn.functional.linear( |
| + experts.act_fn(gate) * up, experts.down_proj[idx] |
| + ) |
| + activations[idx] = out.to(device) |
| + |
| + elif self.hook_config.fused_experts: |
| _, router_scores = output # (num_experts, total_tokens) |
| router_logits = module.router(flat_input) # (total_tokens, num_experts) |
| _, selected_experts = torch.topk(router_logits, top_k, dim=-1) |
| @@ -484,6 +514,22 @@ class Qwen3MoEObserverHookConfig(MoETransformerObserverConfig): |
| module_class_name_to_hook_regex: Optional[str] = "Qwen3MoeSparseMoeBlock" |
| |
| |
| +@dataclass |
| +class Qwen35MoEObserverHookConfig(MoETransformerObserverConfig): |
| + module_class_name_to_hook_regex: Optional[str] = "Qwen3_5MoeSparseMoeBlock" |
| + num_experts_attr_name: str = "experts.num_experts" |
| + top_k_attr_name: str = "gate.top_k" |
| + fused_3d_experts: bool = True |
| + |
| + |
| +@dataclass |
| +class Qwen4ExpMoEObserverHookConfig(MoETransformerObserverConfig): |
| + module_class_name_to_hook_regex: Optional[str] = "Qwen4ExpTextSparseMoeBlock" |
| + num_experts_attr_name: str = "experts.num_experts" |
| + top_k_attr_name: str = "gate.top_k" |
| + fused_3d_experts: bool = True |
| + |
| + |
| @dataclass |
| class Llama4MoEObserverHookConfig(MoETransformerObserverConfig): |
| module_class_name_to_hook_regex: Optional[str] = "Llama4TextMoe" |
| @@ -533,4 +579,6 @@ OBSERVER_CONFIG_REGISTRY = { |
| "Ernie4_5_MoEForCausalLM": Ernie4_5MoEObserverHookConfig, |
| "Ernie4_5_MoeForCausalLM": Ernie4_5MoEObserverHookConfig, |
| "Glm4MoeForCausalLM": Glm44MoEObserverHookConfig, |
| + "Qwen3_5MoeForCausalLM": Qwen35MoEObserverHookConfig, |
| + "Qwen4ExpForCausalLM": Qwen4ExpMoEObserverHookConfig, |
| } |
| |
| |
| |
| |
| @@ -33,7 +33,6 @@ from reap.cluster import ( |
| dynamic_frequency_penalized_clustering, |
| ) |
| from reap.model_util import get_moe, assert_merge, MODEL_ATTRS, patched_model_map, get_super_expert_indices |
| -from reap.eval import run_evaluate |
| import shutil |
| |
| logger = logging.getLogger(__name__) |
| @@ -139,10 +138,23 @@ def prune( |
| ] |
| moe.experts.down_proj.data = moe.experts.down_proj[retained_expert_indicies] |
| moe.num_experts = len(retained_expert_indicies) |
| - moe.router.weight.data = moe.router.weight.data[retained_expert_indicies] |
| - moe.router.out_features = len(retained_expert_indicies) |
| - if hasattr(moe.router, "num_experts"): # transformers >= 4.54+ |
| - moe.router.num_experts = len(retained_expert_indicies) |
| + # apex: the router attribute from MODEL_ATTRS (Llama-4: router; Qwen3.5/4Exp: gate) |
| + router = getattr(moe, model_attrs["router"]) |
| + router.weight.data = router.weight.data[retained_expert_indicies] |
| + router.out_features = len(retained_expert_indicies) |
| + if hasattr(router, "num_experts"): # transformers >= 4.54+ |
| + router.num_experts = len(retained_expert_indicies) |
| + if hasattr(moe.experts, "num_experts"): # apex: transformers 5 3D experts |
| + moe.experts.num_experts = len(retained_expert_indicies) |
| + for name in ( # apex: transformers FP8Experts keep per-expert block scales |
| + "gate_up_proj_scale_inv", |
| + "down_proj_scale_inv", |
| + "gate_up_proj_activation_scale", |
| + "down_proj_activation_scale", |
| + ): |
| + param = getattr(moe.experts, name, None) |
| + if param is not None: |
| + param.data = param.data[retained_expert_indicies] |
| |
| # patch config and dump |
| logger.info("Saving pruned model...") |
| @@ -310,15 +322,15 @@ def main(): |
| |
| logger.info("Pruning completed.") |
| |
| - dump_args_to_yaml( |
| + dump_args_to_yaml( # apex: upstream passed these positionally to a **kwargs signature |
| pruned_model_dir, |
| - reap_args, |
| - ds_args, |
| - obs_args, |
| - model_args, |
| - eval_args, |
| - prune_args, |
| - cluster_args, |
| + reap_args=reap_args, |
| + ds_args=ds_args, |
| + obs_args=obs_args, |
| + model_args=model_args, |
| + eval_args=eval_args, |
| + prune_args=prune_args, |
| + cluster_args=cluster_args, |
| ) |
| |
| # eval |
| @@ -330,6 +342,8 @@ def main(): |
| torch.cuda.empty_cache() |
| gc.collect() |
| model_args.model_name = pruned_model_dir |
| + from reap.eval import run_evaluate # apex: lazy, pulls in vllm / lm_eval / evalplus |
| + |
| run_evaluate(model_args, pruned_model_dir / "eval", eval_args, reap_args.seed) |
| |
| |
|
|
