diff --git a/src/reap/data.py b/src/reap/data.py index c41066d..2831c31 100644 --- a/src/reap/data.py +++ b/src/reap/data.py @@ -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, diff --git a/src/reap/main.py b/src/reap/main.py index acc0a2a..8977325 100644 --- a/src/reap/main.py +++ b/src/reap/main.py @@ -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) diff --git a/src/reap/model_util.py b/src/reap/model_util.py index 16bb5d2..d6a0140 100644 --- a/src/reap/model_util.py +++ b/src/reap/model_util.py @@ -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", diff --git a/src/reap/observer.py b/src/reap/observer.py index 2668200..859021d 100644 --- a/src/reap/observer.py +++ b/src/reap/observer.py @@ -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, } diff --git a/src/reap/prune.py b/src/reap/prune.py index 5767c4e..12d9106 100644 --- a/src/reap/prune.py +++ b/src/reap/prune.py @@ -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)