From b6a5b242cc330ce61663cbe53dd708e7e1eab5f5 Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Sun, 4 Oct 2026 15:29:09 +0000 Subject: [PATCH] chore(models)!: remove GPT-OSS support --- README.md | 1 - docs/advanced.md | 6 +- .../src/prime_rl/utils/parsers.py | 2 - src/prime_rl/inference/patches.py | 2 - .../inference/vllm/gpt_oss_weight_loading.py | 145 ---------- src/prime_rl/trainer/lora.py | 8 +- src/prime_rl/trainer/models/__init__.py | 3 - .../trainer/models/gpt_oss/__init__.py | 13 - .../trainer/models/gpt_oss/attention.py | 181 ------------- .../models/gpt_oss/configuration_gpt_oss.py | 112 -------- .../models/gpt_oss/converting_gpt_oss.py | 64 ----- .../models/gpt_oss/modeling_gpt_oss.py | 227 ---------------- .../trainer/models/layers/activations.py | 12 +- src/prime_rl/trainer/models/layers/attn.py | 4 - .../trainer/models/layers/expert_compute.py | 28 +- .../trainer/models/layers/lora/multi_moe.py | 256 ------------------ src/prime_rl/trainer/models/layers/moe.py | 32 +-- .../trainer/models/layers/sonic_moe.py | 6 +- .../trainer/models/layers/ulysses_attn.py | 12 - tests/unit/test_parsers.py | 3 - tests/unit/train/models/test_gpt_oss.py | 133 --------- tests/unit/train/models/test_moe.py | 16 +- 22 files changed, 14 insertions(+), 1252 deletions(-) delete mode 100644 src/prime_rl/inference/vllm/gpt_oss_weight_loading.py delete mode 100644 src/prime_rl/trainer/models/gpt_oss/__init__.py delete mode 100644 src/prime_rl/trainer/models/gpt_oss/attention.py delete mode 100644 src/prime_rl/trainer/models/gpt_oss/configuration_gpt_oss.py delete mode 100644 src/prime_rl/trainer/models/gpt_oss/converting_gpt_oss.py delete mode 100644 src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py delete mode 100644 tests/unit/train/models/test_gpt_oss.py diff --git a/README.md b/README.md index 666c8af0195..16086fbb8b1 100644 --- a/README.md +++ b/README.md @@ -57,7 +57,6 @@ The trainer runs the PrimeRL model implementations under `src/prime_rl/trainer/m | Nemotron H (`nemotron_h`) | `nvidia/Nemotron-3-Nano-30B-A3B`, `nvidia/Nemotron-3-Super-120B-A12B`, … | yes | ✅ | ✅ | | Trinity (`afmoe`) | `arcee-ai/Trinity-Mini`, … | yes | ✅ | ✅ | | GLM-4 · GLM-4.5 MoE · INTELLECT-3 (`glm4_moe`) | `THUDM/GLM-4-9B-0414`, `zai-org/GLM-4.5-Air`, `zai-org/GLM-4.5`, `PrimeIntellect/INTELLECT-3`, … | yes | ✅ | ✅ | -| GPT-OSS (`gpt_oss`) | `unsloth/gpt-oss-20b-BF16`, … | yes | ✅ | ✅ | ## Setup diff --git a/docs/advanced.md b/docs/advanced.md index ec6d7c43abd..eccfa4be19e 100644 --- a/docs/advanced.md +++ b/docs/advanced.md @@ -29,12 +29,8 @@ The trainer only runs `prime-rl`'s own model implementations, selected from the | Nemotron H | `nvidia/Nemotron-3-Nano-30B-A3B`, … | ✅ | ❌ | | Trinity (AFMoE) | `arcee-ai/Trinity-Mini`, … | ✅ | ✅ | | GLM-4 / GLM-4.5 / INTELLECT-3 | `THUDM/GLM-4-9B-0414`, `zai-org/GLM-4.5`, `PrimeIntellect/INTELLECT-3`, … | ✅ | ✅ | -| GPT-OSS | `unsloth/gpt-oss-20b-BF16`, … | ✅ | ✅ | | DeepSeek V4 | `deepseek-ai/DeepSeek-V4-Flash-0731` | ✅ | ✅ | -GPT-OSS uses FlashAttention 4 with learned attention sinks. Training requires SM90 or SM100/SM110 GPUs -and a BF16 checkpoint such as `unsloth/gpt-oss-20b-BF16`; the original MXFP4 checkpoints are not supported. - ### Low-precision training Dense linear precision and routed-expert precision are configured independently. `[trainer.model.quantization]` applies only to dense `Linear` modules: @@ -48,7 +44,7 @@ Dense linear precision and routed-expert precision are configured independently. - `type = "deepgemm_fp8"` (requires DeepGEMM and SM90+) - `type = "mxfp8"` (requires `prime-kernels`, torchao, and SM100) -SonicMoE uses the upstream `sonic-moe` package (`uv sync --extra sonic-moe`) for fused BF16 expert computation. The supported model is Qwen3 MoE with the `gate_up` model fusion enabled. Backend selection requires fused gate/up weights, standard SwiGLU, and bias-free experts; incompatible expert structures raise an error during setup. It uses the same router and local, torch EP, or DeepEP dispatch as other compute backends: +SonicMoE uses the upstream `sonic-moe` package (`uv sync --extra sonic-moe`) for fused BF16 expert computation. The supported model is Qwen3 MoE with the `gate_up` model fusion enabled. Backend selection requires fused gate/up weights and standard SwiGLU; incompatible expert structures raise an error during setup. It uses the same router and local, torch EP, or DeepEP dispatch as other compute backends: ```toml [trainer.model] diff --git a/packages/prime-rl-configs/src/prime_rl/utils/parsers.py b/packages/prime-rl-configs/src/prime_rl/utils/parsers.py index 7ff04b99647..44d7a7fcd0f 100644 --- a/packages/prime-rl-configs/src/prime_rl/utils/parsers.py +++ b/packages/prime-rl-configs/src/prime_rl/utils/parsers.py @@ -12,7 +12,6 @@ (re.compile(r"^zai-org/GLM-4\.7"), "glm47"), (re.compile(r"^zai-org/GLM-5"), "glm47"), (re.compile(r"^MiniMaxAI/MiniMax-M2"), "minimax_m2"), - (re.compile(r"^openai/gpt-oss"), "openai"), (re.compile(r"^poolside/Laguna"), "poolside_v1"), (re.compile(r"^PrimeIntellect/INTELLECT-3"), "qwen3_coder"), (re.compile(r"^nvidia/NVIDIA-Nemotron-3"), "qwen3_coder"), @@ -32,7 +31,6 @@ (re.compile(r"^deepseek-ai/DeepSeek-V3\.[12]"), "deepseek_r1"), (re.compile(r"^zai-org/GLM-"), "glm45"), (re.compile(r"^MiniMaxAI/MiniMax-M2"), "minimax_m2_append_think"), - # openai/gpt-oss has no entry here — vLLM's harmony serving path splits reasoning natively. (re.compile(r"^poolside/Laguna"), "poolside_v1"), (re.compile(r"^PrimeIntellect/INTELLECT-3"), "deepseek_r1"), (re.compile(r"^nvidia/NVIDIA-Nemotron-3-Super"), "nemotron_v3"), diff --git a/src/prime_rl/inference/patches.py b/src/prime_rl/inference/patches.py index 972faf4fef8..26b5c4b2bd9 100644 --- a/src/prime_rl/inference/patches.py +++ b/src/prime_rl/inference/patches.py @@ -9,10 +9,8 @@ def apply_shared_vllm_patches(): vLLM swallows plugin load failures (``load_plugins_by_group`` logs and continues), so a broken entry-point target silently skips ALL of these patches. """ - from prime_rl.inference.vllm.gpt_oss_weight_loading import patch_gpt_oss_weight_loading from prime_rl.inference.vllm.qwen38_weight_loading import patch_qwen38_weight_loading - patch_gpt_oss_weight_loading() patch_qwen38_weight_loading() _patch_lora_key_prefix() _patch_qwen35_moe_lora_format() diff --git a/src/prime_rl/inference/vllm/gpt_oss_weight_loading.py b/src/prime_rl/inference/vllm/gpt_oss_weight_loading.py deleted file mode 100644 index 247c516f4c9..00000000000 --- a/src/prime_rl/inference/vllm/gpt_oss_weight_loading.py +++ /dev/null @@ -1,145 +0,0 @@ -from collections.abc import Iterable - -import torch - - -def patch_gpt_oss_weight_loading() -> None: - """Route GPT-OSS BF16 expert and sink weights through vLLM loaders. - - Remove this patch when the pinned vLLM version uses weight loaders for these - prepared tensors in ``GptOssModel._load_weights_other``. - """ - from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts - from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, - maybe_remap_moe_expert_param_name, - ) - from vllm.model_executor.models import gpt_oss - - original_load_weights = gpt_oss.GptOssModel._load_weights_other - if getattr(original_load_weights, "_prime_rl_uses_weight_loaders", False): - return - - original_expert_weight_loader = RoutedExperts.weight_loader - - def expert_weight_loader( - self, - param, - loaded_weight, - weight_name, - shard_id, - expert_id, - return_success=False, - ): - if expert_id is None: - if shard_id is not None: - raise ValueError("A combined expert tensor must already be sharded") - if param.shape != loaded_weight.shape: - raise ValueError( - f"Combined expert tensor shape {tuple(loaded_weight.shape)} does not match " - f"parameter shape {tuple(param.shape)}" - ) - param.copy_(loaded_weight) - return True if return_success else None - return original_expert_weight_loader( - self, - param, - loaded_weight, - weight_name, - shard_id, - expert_id, - return_success, - ) - - RoutedExperts.weight_loader = expert_weight_loader - - def load_weights( - self, - ep_rank_end: int, - ep_rank_start: int, - heads_per_rank: int, - head_start: int, - weights: Iterable[tuple[str, torch.Tensor]], - stacked_params_mapping: list[tuple[str, ...]], - ) -> set[str]: - params = dict(self.named_parameters()) - loaded_params: set[str] = set() - use_ep = self.parallel_config.enable_expert_parallel - tp_size, tp_rank = gpt_oss.FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp( - tp_size=gpt_oss.get_tensor_model_parallel_world_size(), - dp_size=gpt_oss.get_dp_group().world_size, - dp_rank=gpt_oss.get_dp_group().rank_in_group, - pcp_size=gpt_oss.get_pcp_group().world_size, - pcp_rank=gpt_oss.get_pcp_group().rank_in_group, - ) - intermediate_size = self.config.intermediate_size - per_rank_intermediate_size = gpt_oss.cdiv(intermediate_size, tp_size) - tp_rank_start = tp_rank * per_rank_intermediate_size - tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size, intermediate_size) - - def remaining_weights(): - for name, weight in weights: - if not ( - name.endswith((".w13_weight", ".w2_weight", ".w13_bias", ".w2_bias")) or name.endswith(".sinks") - ): - yield name, weight - continue - - name = maybe_remap_moe_expert_param_name(name, params) - if gpt_oss.is_pp_missing_parameter(name, self): - continue - - param = params[name] - if name.endswith(".w13_weight"): - if use_ep: - weight = weight[ep_rank_start:ep_rank_end] - else: - weight = weight[:, :, 2 * tp_rank_start : 2 * tp_rank_end] - weight = weight.permute(0, 2, 1).contiguous() - elif name.endswith(".w2_weight"): - if use_ep: - weight = weight[ep_rank_start:ep_rank_end] - else: - weight = weight[:, tp_rank_start:tp_rank_end] - weight = weight.permute(0, 2, 1).contiguous() - elif name.endswith(".w13_bias"): - if use_ep: - weight = weight[ep_rank_start:ep_rank_end] - else: - weight = weight[:, 2 * tp_rank_start : 2 * tp_rank_end] - elif name.endswith(".w2_bias"): - if use_ep: - weight = weight[ep_rank_start:ep_rank_end] - elif tp_rank != 0: - weight = weight.zero_() - else: - weight = weight.narrow(0, head_start, heads_per_rank) - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, weight) - loaded_params.add(name) - continue - - param.weight_loader( - param, - weight, - weight_name=name, - shard_id=None, - expert_id=None, - ) - loaded_params.add(name) - - loaded_params.update( - original_load_weights( - self, - ep_rank_end, - ep_rank_start, - heads_per_rank, - head_start, - remaining_weights(), - stacked_params_mapping, - ) - ) - return loaded_params - - load_weights._prime_rl_uses_weight_loaders = True - gpt_oss.GptOssModel._load_weights_other = load_weights diff --git a/src/prime_rl/trainer/lora.py b/src/prime_rl/trainer/lora.py index 0a48257b3db..a529c6d4795 100644 --- a/src/prime_rl/trainer/lora.py +++ b/src/prime_rl/trainer/lora.py @@ -14,7 +14,6 @@ set_multilora_scaling, ) from prime_rl.trainer.models.layers.lora.multi_moe import ( - MultiLoRAGptOssGroupedExperts, MultiLoRAGroupedExperts, MultiLoRANonGatedGroupedExperts, ) @@ -224,9 +223,6 @@ def apply_lora_to_model(model: nn.Module, config: LoRAConfig) -> None: lora_state = setup_lora_state(config, torch.device("cuda", get_world().local_rank)) if isinstance(model, PreTrainedModelPrimeRL): lora_state.register_adapter_state_dict_converter(type(model).convert_adapter_to_hf) - uses_gpt_oss_moe_adapter = ( - isinstance(model, PreTrainedModelPrimeRL) and getattr(model.config, "model_type", None) == "gpt_oss" - ) from torch.distributed.fsdp import FSDPModule @@ -259,9 +255,7 @@ def apply_lora_to_model(model: nn.Module, config: LoRAConfig) -> None: ) # Handle GroupedExperts (MoE) elif isinstance(base_module, GroupedExperts): - if uses_gpt_oss_moe_adapter: - wrapper = MultiLoRAGptOssGroupedExperts - elif base_module.gate_proj is not None: + if base_module.gate_proj is not None: wrapper = MultiLoRAGroupedExperts else: wrapper = MultiLoRANonGatedGroupedExperts diff --git a/src/prime_rl/trainer/models/__init__.py b/src/prime_rl/trainer/models/__init__.py index af9a0684f14..9d6749a5b0a 100644 --- a/src/prime_rl/trainer/models/__init__.py +++ b/src/prime_rl/trainer/models/__init__.py @@ -14,7 +14,6 @@ from prime_rl.trainer.models.deepseek_v4 import DeepseekV4Config, DeepseekV4ForCausalLM from prime_rl.trainer.models.glm4_moe import Glm4MoeConfig, Glm4MoeForCausalLM from prime_rl.trainer.models.glm_moe_dsa import GlmMoeDsaConfig, GlmMoeDsaForCausalLM -from prime_rl.trainer.models.gpt_oss import GptOssConfig, GptOssForCausalLM from prime_rl.trainer.models.laguna import LagunaConfig, LagunaForCausalLM from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, cast_float_and_contiguous from prime_rl.trainer.models.llama import LlamaForCausalLM @@ -41,7 +40,6 @@ AutoConfig.register("deepseek_v4", DeepseekV4Config, exist_ok=True) AutoConfig.register("glm4_moe", Glm4MoeConfig, exist_ok=True) AutoConfig.register("glm_moe_dsa", GlmMoeDsaConfig, exist_ok=True) -AutoConfig.register("gpt_oss", GptOssConfig, exist_ok=True) AutoConfig.register("laguna", LagunaConfig, exist_ok=True) AutoConfig.register("minimax_m2", MiniMaxM2Config, exist_ok=True) AutoConfig.register("nemotron_h", NemotronHConfig, exist_ok=True) @@ -75,7 +73,6 @@ (Qwen3_5MoeConfig, Qwen3_5ForCausalLM), (Qwen3_8FlashNextTextConfig, Qwen3_8FlashNextForCausalLM), (Qwen3_8FlashNextConfig, Qwen3_8FlashNextForCausalLM), - (GptOssConfig, GptOssForCausalLM), ) for config_cls, model_cls in _CUSTOM_CAUSAL_LM_MODELS: _CUSTOM_CAUSAL_LM_MAPPING.register(config_cls, model_cls, exist_ok=True) diff --git a/src/prime_rl/trainer/models/gpt_oss/__init__.py b/src/prime_rl/trainer/models/gpt_oss/__init__.py deleted file mode 100644 index 6c72b3a1d7f..00000000000 --- a/src/prime_rl/trainer/models/gpt_oss/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -from prime_rl.trainer.models.gpt_oss.configuration_gpt_oss import GptOssConfig -from prime_rl.trainer.models.gpt_oss.modeling_gpt_oss import ( - GptOssForCausalLM, - GptOssModel, - GptOssPreTrainedModel, -) - -__all__ = [ - "GptOssConfig", - "GptOssForCausalLM", - "GptOssModel", - "GptOssPreTrainedModel", -] diff --git a/src/prime_rl/trainer/models/gpt_oss/attention.py b/src/prime_rl/trainer/models/gpt_oss/attention.py deleted file mode 100644 index 73cf2a6db2a..00000000000 --- a/src/prime_rl/trainer/models/gpt_oss/attention.py +++ /dev/null @@ -1,181 +0,0 @@ -"""GPT-OSS attention and context-parallel execution. - -Learned sink logits are model parameters, not an attention-backend option. GPT-OSS -therefore owns its FA4 call and shards each sink with the corresponding query heads -in its ring and Ulysses paths. -""" - -import torch -from torch import nn - -from prime_rl.trainer.models.fusions import fuse_qkv_projections -from prime_rl.trainer.models.gpt_oss.configuration_gpt_oss import GptOssConfig -from prime_rl.trainer.models.layers.attn import flash_attn_4_varlen_func -from prime_rl.trainer.models.layers.rotary_emb import apply_rotary_pos_emb - - -class GptOssAttention(nn.Module): - supported_fusions = {"qkv": fuse_qkv_projections} - - def __init__(self, config: GptOssConfig, layer_idx: int) -> None: - super().__init__() - self.head_dim = config.head_dim - self.num_attention_heads = config.num_attention_heads - self.num_key_value_heads = config.num_key_value_heads - self.scaling = self.head_dim**-0.5 - self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None - - self.q_proj = nn.Linear( - config.hidden_size, - config.num_attention_heads * config.head_dim, - bias=config.attention_bias, - ) - self.k_proj = nn.Linear( - config.hidden_size, - config.num_key_value_heads * config.head_dim, - bias=config.attention_bias, - ) - self.v_proj = nn.Linear( - config.hidden_size, - config.num_key_value_heads * config.head_dim, - bias=config.attention_bias, - ) - self.qkv_sizes = (self.q_proj.out_features, self.k_proj.out_features, self.v_proj.out_features) - self.register_module("qkv_proj", None) - self.o_proj = nn.Linear( - config.num_attention_heads * config.head_dim, - config.hidden_size, - bias=config.attention_bias, - ) - self.sinks = nn.Parameter(torch.empty(config.num_attention_heads)) - nn.init.normal_(self.sinks, mean=0.0, std=config.initializer_range) - self.flash_attn = flash_attn_4_varlen_func - - def compute_attention( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - cu_seqlens: torch.Tensor, - max_seqlen: int, - ) -> torch.Tensor: - window_size = (self.sliding_window - 1, 0) if self.sliding_window is not None else (None, None) - output, _ = self.flash_attn( - query, - key, - value, - cu_seqlens_q=cu_seqlens, - cu_seqlens_k=cu_seqlens, - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, - softmax_scale=self.scaling, - causal=True, - window_size=window_size, - learnable_sink=self.sinks, - ) - return output - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - cu_seqlens: torch.Tensor, - max_seqlen: int, - ) -> torch.Tensor: - batch_size, sequence_length, _ = hidden_states.shape - if self.qkv_proj is None: - query, key, value = self.q_proj(hidden_states), self.k_proj(hidden_states), self.v_proj(hidden_states) - else: - query, key, value = self.qkv_proj(hidden_states).split(self.qkv_sizes, dim=-1) - query = query.view(batch_size, sequence_length, self.num_attention_heads, self.head_dim) - key = key.view(batch_size, sequence_length, self.num_key_value_heads, self.head_dim) - value = value.view(batch_size, sequence_length, self.num_key_value_heads, self.head_dim) - - cos, sin = position_embeddings - query, key = apply_rotary_pos_emb( - query.transpose(1, 2), - key.transpose(1, 2), - cos, - sin, - ) - query = query.transpose(1, 2) - key = key.transpose(1, 2) - - output = self.compute_attention(query[0], key[0], value[0], cu_seqlens, max_seqlen) - return self.o_proj(output.contiguous().view(batch_size, sequence_length, -1)) - - -def substitute_gpt_oss_ring_attention( - process_group: torch.distributed.ProcessGroup, - heads_k_stride: int, -) -> None: - from ring_flash_attn.adapters.hf_adapter import DATA_PARAMS - - from prime_rl.trainer.distributed.collectives import all_gather - - def ring_attention(self, query, key, value, _cu_seqlens, _max_seqlen): - key_value_groups = self.num_attention_heads // self.num_key_value_heads - local_k_slice = DATA_PARAMS["local_k_slice"] - window_size = (self.sliding_window - 1, 0) if self.sliding_window is not None else (None, None) - outputs = [] - - for key_head_start in range(0, self.num_key_value_heads, heads_k_stride): - key_head_stop = min(key_head_start + heads_k_stride, self.num_key_value_heads) - query_head_start = key_head_start * key_value_groups - query_head_stop = key_head_stop * key_value_groups - - gathered_key = all_gather(key[:, key_head_start:key_head_stop], 0, process_group)[local_k_slice] - gathered_value = all_gather(value[:, key_head_start:key_head_stop], 0, process_group)[local_k_slice] - output, _ = self.flash_attn( - query[:, query_head_start:query_head_stop], - gathered_key, - gathered_value, - cu_seqlens_q=DATA_PARAMS["cu_seqlens_q"], - cu_seqlens_k=DATA_PARAMS["cu_seqlens_k"], - max_seqlen_q=DATA_PARAMS["max_seqlen_q"], - max_seqlen_k=DATA_PARAMS["max_seqlen_k"], - softmax_scale=self.scaling, - causal=True, - window_size=window_size, - learnable_sink=self.sinks[query_head_start:query_head_stop], - ) - outputs.append(output) - - return torch.cat(outputs, dim=1) - - GptOssAttention.compute_attention = ring_attention - - -def substitute_gpt_oss_ulysses_attention(process_group: torch.distributed.ProcessGroup) -> None: - from prime_rl.trainer.models.layers.ulysses_attn import ULYSSES_PARAMS, ulysses_flash_attn_varlen_func - - cp_size = torch.distributed.get_world_size(process_group) - - def ulysses_attention(self, query, key, value, _cu_seqlens, _max_seqlen): - window_size = (self.sliding_window - 1, 0) if self.sliding_window is not None else (-1, -1) - return ulysses_flash_attn_varlen_func( - self.flash_attn, - query, - key, - value, - cu_seqlens_q=ULYSSES_PARAMS["cu_seqlens"], - cu_seqlens_k=ULYSSES_PARAMS["cu_seqlens"], - max_seqlen_q=ULYSSES_PARAMS["max_seqlen"], - max_seqlen_k=ULYSSES_PARAMS["max_seqlen"], - causal=True, - cp_group=process_group, - cp_size=cp_size, - flash_attn_version=4, - window_size=window_size, - softmax_scale=self.scaling, - learnable_sink=self.sinks, - ) - - GptOssAttention.compute_attention = ulysses_attention - - -__all__ = [ - "GptOssAttention", - "substitute_gpt_oss_ring_attention", - "substitute_gpt_oss_ulysses_attention", -] diff --git a/src/prime_rl/trainer/models/gpt_oss/configuration_gpt_oss.py b/src/prime_rl/trainer/models/gpt_oss/configuration_gpt_oss.py deleted file mode 100644 index f174e2f9be5..00000000000 --- a/src/prime_rl/trainer/models/gpt_oss/configuration_gpt_oss.py +++ /dev/null @@ -1,112 +0,0 @@ -from transformers.configuration_utils import PretrainedConfig - - -class GptOssConfig(PretrainedConfig): - model_type = "gpt_oss" - attribute_map = {"num_experts": "num_local_experts"} - - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.experts.gate_proj": "colwise", - "layers.*.mlp.experts.up_proj": "colwise", - "layers.*.mlp.experts.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - base_model_ep_plan = { - "layers.*.mlp.router": "ep_router", - "layers.*.mlp.experts.gate_proj": "grouped_gemm", - "layers.*.mlp.experts.gate_proj_bias": "grouped_gemm", - "layers.*.mlp.experts.up_proj": "grouped_gemm", - "layers.*.mlp.experts.up_proj_bias": "grouped_gemm", - "layers.*.mlp.experts.down_proj": "grouped_gemm", - "layers.*.mlp.experts.down_proj_bias": "grouped_gemm", - "layers.*.mlp.experts": "moe_tp_experts", - } - - def __init__( - self, - vocab_size: int = 201088, - hidden_size: int = 2880, - intermediate_size: int = 2880, - num_hidden_layers: int = 36, - num_attention_heads: int = 64, - num_key_value_heads: int = 8, - head_dim: int = 64, - num_local_experts: int = 128, - num_experts_per_tok: int = 4, - hidden_act: str = "silu", - swiglu_limit: float = 7.0, - max_position_embeddings: int = 131072, - initial_context_length: int = 4096, - sliding_window: int | None = 128, - layer_types: list[str] | None = None, - rope_theta: float = 150000.0, - rope_scaling: dict | None = None, - rope_parameters: dict | None = None, - rms_norm_eps: float = 1e-5, - attention_bias: bool = True, - attention_dropout: float = 0.0, - initializer_range: float = 0.02, - router_aux_loss_coef: float = 0.001, - output_router_logits: bool = False, - use_cache: bool = True, - pad_token_id: int | None = None, - bos_token_id: int | None = None, - eos_token_id: int | list[int] | None = None, - tie_word_embeddings: bool = False, - **kwargs, - ) -> None: - self.vocab_size = vocab_size - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.num_key_value_heads = num_key_value_heads - self.head_dim = head_dim - self.num_local_experts = num_local_experts - self.num_experts_per_tok = num_experts_per_tok - self.hidden_act = hidden_act - self.swiglu_limit = swiglu_limit - self.max_position_embeddings = max_position_embeddings - self.initial_context_length = initial_context_length - self.sliding_window = sliding_window - self.layer_types = layer_types or [ - "sliding_attention" if layer_idx % 2 == 0 else "full_attention" for layer_idx in range(num_hidden_layers) - ] - - if rope_parameters is None: - rope_parameters = rope_scaling - if rope_parameters is None: - rope_parameters = { - "rope_type": "yarn", - "factor": 32.0, - "beta_fast": 32.0, - "beta_slow": 1.0, - "truncate": False, - "original_max_position_embeddings": 4096, - } - self.rope_parameters = dict(rope_parameters) - self.rope_parameters.setdefault("rope_theta", rope_theta) - - self.rms_norm_eps = rms_norm_eps - self.attention_bias = attention_bias - self.attention_dropout = attention_dropout - self.initializer_range = initializer_range - self.router_aux_loss_coef = router_aux_loss_coef - self.output_router_logits = output_router_logits - self.use_cache = use_cache - self.pad_token_id = pad_token_id - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id - - super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) - - -__all__ = ["GptOssConfig"] diff --git a/src/prime_rl/trainer/models/gpt_oss/converting_gpt_oss.py b/src/prime_rl/trainer/models/gpt_oss/converting_gpt_oss.py deleted file mode 100644 index d1b5815320a..00000000000 --- a/src/prime_rl/trainer/models/gpt_oss/converting_gpt_oss.py +++ /dev/null @@ -1,64 +0,0 @@ -"""HF <-> PrimeRL weight conversion for GPT-OSS.""" - -from torch import Tensor - -from prime_rl.trainer.models.conversion_ops import ConvOp, Rename - - -class GptOssExperts(ConvOp): - def __init__(self, prefix: str) -> None: - self.prefix = prefix - - def hf_to_prime(self, state_dict: dict[str, Tensor]) -> None: - gate_up_name = f"{self.prefix}.gate_up_proj" - if gate_up_name not in state_dict: - return - - gate_up = state_dict.pop(gate_up_name) - gate_up_bias = state_dict.pop(f"{self.prefix}.gate_up_proj_bias") - state_dict[f"{self.prefix}.gate_proj"] = gate_up[..., ::2].transpose(-2, -1) - state_dict[f"{self.prefix}.up_proj"] = gate_up[..., 1::2].transpose(-2, -1) - state_dict[f"{self.prefix}.gate_proj_bias"] = gate_up_bias[..., ::2] - state_dict[f"{self.prefix}.up_proj_bias"] = gate_up_bias[..., 1::2] - state_dict[f"{self.prefix}.down_proj"] = state_dict[f"{self.prefix}.down_proj"].transpose(-2, -1) - - def prime_to_hf(self, state_dict: dict[str, Tensor]) -> None: - gate_name = f"{self.prefix}.gate_proj" - if gate_name not in state_dict: - return - - gate = state_dict.pop(gate_name).transpose(-2, -1) - up = state_dict.pop(f"{self.prefix}.up_proj").transpose(-2, -1) - gate_bias = state_dict.pop(f"{self.prefix}.gate_proj_bias") - up_bias = state_dict.pop(f"{self.prefix}.up_proj_bias") - state_dict[f"{self.prefix}.gate_up_proj"] = gate.new_empty((*gate.shape[:-1], 2 * gate.shape[-1])) - state_dict[f"{self.prefix}.gate_up_proj"][..., ::2] = gate - state_dict[f"{self.prefix}.gate_up_proj"][..., 1::2] = up - state_dict[f"{self.prefix}.gate_up_proj_bias"] = gate_bias.new_empty( - (*gate_bias.shape[:-1], 2 * gate_bias.shape[-1]) - ) - state_dict[f"{self.prefix}.gate_up_proj_bias"][..., ::2] = gate_bias - state_dict[f"{self.prefix}.gate_up_proj_bias"][..., 1::2] = up_bias - state_dict[f"{self.prefix}.down_proj"] = state_dict[f"{self.prefix}.down_proj"].transpose(-2, -1) - - -def is_hf_state_dict(state_dict: dict[str, Tensor]) -> bool: - return any(name.endswith("mlp.experts.gate_up_proj") for name in state_dict) - - -def is_prime_state_dict(state_dict: dict[str, Tensor]) -> bool: - return any("mlp.experts.gate_proj" in name for name in state_dict) - - -def conversion_chain(config) -> list[ConvOp]: - operations: list[ConvOp] = [] - for layer_idx in range(config.num_hidden_layers): - prefix = f"model.layers.{layer_idx}.mlp" - operations.extend( - [ - Rename(f"{prefix}.router.weight", f"{prefix}.router.gate.weight"), - Rename(f"{prefix}.router.bias", f"{prefix}.router.gate.bias"), - GptOssExperts(f"{prefix}.experts"), - ] - ) - return operations diff --git a/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py b/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py deleted file mode 100644 index 36a60c780e3..00000000000 --- a/src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py +++ /dev/null @@ -1,227 +0,0 @@ -import torch -from torch import Tensor, nn -from transformers.modeling_outputs import BaseModelOutput - -from prime_rl.trainer.models.base import PreTrainedModelPrimeRL -from prime_rl.trainer.models.gpt_oss.attention import GptOssAttention -from prime_rl.trainer.models.gpt_oss.configuration_gpt_oss import GptOssConfig -from prime_rl.trainer.models.gpt_oss.converting_gpt_oss import ( - conversion_chain, - is_hf_state_dict, - is_prime_state_dict, -) -from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, VanillaOutputLinear -from prime_rl.trainer.models.layers.moe import GroupedExperts, MoE, TokenChoiceTopKRouter -from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig -from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens - - -class GptOssRMSNorm(nn.Module): - def __init__(self, hidden_size: int, eps: float) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - input_dtype = hidden_states.dtype - hidden_states = hidden_states.float() - variance = hidden_states.square().mean(dim=-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return (self.weight * hidden_states).to(input_dtype) - - -class GptOssDecoderLayer(nn.Module): - def __init__(self, config: GptOssConfig, layer_idx: int) -> None: - super().__init__() - self.self_attn = GptOssAttention(config, layer_idx) - - router = TokenChoiceTopKRouter( - dim=config.hidden_size, - num_experts=config.num_local_experts, - top_k=config.num_experts_per_tok, - score_func="topk_softmax", - route_norm=False, - route_scale=1.0, - gate_bias=True, - ) - experts = GroupedExperts( - dim=config.hidden_size, - hidden_dim=config.intermediate_size, - num_experts=config.num_local_experts, - expert_type="gated", - activation="clamped_swiglu", - bias=True, - ) - experts.init_weights(config.initializer_range) - self.mlp = MoE( - router=router, - experts=experts, - shared_expert=None, - score_before_experts=False, - load_balance_coeff=None, - ) - self.input_layernorm = GptOssRMSNorm(config.hidden_size, config.rms_norm_eps) - self.post_attention_layernorm = GptOssRMSNorm(config.hidden_size, config.rms_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - cu_seqlens: torch.Tensor, - max_seqlen: int, - routed_experts: torch.Tensor | None = None, - ) -> torch.Tensor: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn(hidden_states, position_embeddings, cu_seqlens, max_seqlen) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states, routed_experts=routed_experts) - return residual + hidden_states - - -class GptOssPreTrainedModel(PreTrainedModelPrimeRL): - config: GptOssConfig - base_model_prefix = "model" - supports_gradient_checkpointing = True - _no_split_modules = ["GptOssDecoderLayer"] - _supports_flash_attn = True - _supports_sdpa = False - _can_compile_fullgraph = False - _supports_attention_backend = True - _keep_in_fp32_modules = ["post_attention_layernorm", "input_layernorm", "norm"] - - @classmethod - def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: - return is_hf_state_dict(state_dict) - - @classmethod - def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: - return is_prime_state_dict(state_dict) - - @classmethod - def conversion_chain(cls, config: GptOssConfig): - return conversion_chain(config) - - -class GptOssModel(GptOssPreTrainedModel): - def __init__(self, config: GptOssConfig) -> None: - super().__init__(config) - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) - self.layers = nn.ModuleList( - GptOssDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers) - ) - self.norm = GptOssRMSNorm(config.hidden_size, config.rms_norm_eps) - self.rotary_emb = RotaryEmbedding( - RotaryEmbeddingConfig( - max_position_embeddings=config.max_position_embeddings, - rope_type=config.rope_parameters["rope_type"], - model_config=config, - ) - ) - self.gradient_checkpointing = False - self.post_init() - - def forward( - self, - input_ids: torch.LongTensor | None = None, - position_ids: torch.LongTensor | None = None, - inputs_embeds: torch.Tensor | None = None, - routed_experts: torch.LongTensor | None = None, - *, - seq_lens: torch.LongTensor, - seq_lens_are_pre_shard: bool = False, - ) -> BaseModelOutput: - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens( - seq_lens.to(device=inputs_embeds.device), - total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1], - ) - torch._dynamo.mark_dynamic(cu_seqlens, 0) - - hidden_states = inputs_embeds - position_embeddings = self.rotary_emb(hidden_states, position_ids) - for layer_idx, decoder_layer in enumerate(self.layers): - layer_routed_experts = routed_experts[:, :, layer_idx] if routed_experts is not None else None - hidden_states = decoder_layer( - hidden_states, - position_embeddings, - cu_seqlens, - max_seqlen, - routed_experts=layer_routed_experts, - ) - return BaseModelOutput(last_hidden_state=self.norm(hidden_states)) - - -class GptOssForCausalLM(GptOssPreTrainedModel): - _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - _tp_plan = {"lm_head": "colwise_gather_output"} - _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} - - def __init__(self, config: GptOssConfig) -> None: - super().__init__(config) - self.model = GptOssModel(config) - self.vocab_size = config.vocab_size - self.lm_head = VanillaOutputLinear(config.hidden_size, config.vocab_size) - self.num_experts = config.num_local_experts - self.num_experts_per_tok = config.num_experts_per_tok - self.post_init() - - def forward( - self, - input_ids: torch.LongTensor | None = None, - position_ids: torch.LongTensor | None = None, - inputs_embeds: torch.Tensor | None = None, - labels: torch.LongTensor | None = None, - logits_to_keep: int | torch.Tensor = 0, - temperature: torch.Tensor | None = None, - routed_experts: torch.LongTensor | None = None, - *, - seq_lens: torch.LongTensor, - seq_lens_are_pre_shard: bool = False, - ) -> PrimeLmOutput: - outputs = self.model( - input_ids=input_ids, - position_ids=position_ids, - inputs_embeds=inputs_embeds, - routed_experts=routed_experts, - seq_lens=seq_lens, - seq_lens_are_pre_shard=seq_lens_are_pre_shard, - ) - hidden_states = outputs.last_hidden_state - if isinstance(logits_to_keep, int): - slice_indices = slice(-logits_to_keep, None) if logits_to_keep > 0 else slice(None) - else: - slice_indices = logits_to_keep - return self.lm_head( - hidden_states[:, slice_indices], - labels[:, slice_indices] if labels is not None else None, - temperature=temperature[:, slice_indices] if temperature is not None else None, - ) - - def init_buffers_post_meta(self) -> None: - rotary_emb = self.model.rotary_emb - inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn( - rotary_emb.config, - rotary_emb.inv_freq.device, - ) - rotary_emb.inv_freq.copy_(inv_freq) - for module in self.modules(): - if isinstance(module, MoE): - module.tokens_per_expert.zero_() - module.routing_confidence_sum.zero_() - - -__all__ = [ - "GptOssDecoderLayer", - "GptOssForCausalLM", - "GptOssModel", - "GptOssPreTrainedModel", - "GptOssRMSNorm", -] diff --git a/src/prime_rl/trainer/models/layers/activations.py b/src/prime_rl/trainer/models/layers/activations.py index a6aa039aa69..a3fa73ea78b 100644 --- a/src/prime_rl/trainer/models/layers/activations.py +++ b/src/prime_rl/trainer/models/layers/activations.py @@ -9,15 +9,6 @@ class Activation(Protocol): def apply(gate: torch.Tensor | None, up: torch.Tensor) -> torch.Tensor: ... -class ClampedSwiglu(Activation): - @staticmethod - def apply(gate: torch.Tensor | None, up: torch.Tensor) -> torch.Tensor: - assert gate is not None - gate = gate.clamp(max=7.0) - up = up.clamp(min=-7.0, max=7.0) - return (up + 1) * gate * torch.sigmoid(gate * 1.702) - - class Relu2(Activation): @staticmethod def apply(gate: torch.Tensor | None, up: torch.Tensor) -> torch.Tensor: @@ -34,10 +25,9 @@ def apply(gate: torch.Tensor | None, up: torch.Tensor) -> torch.Tensor: return F.silu(gate) * up -ActivationType = Literal["silu", "relu2", "clamped_swiglu"] +ActivationType = Literal["silu", "relu2"] ActivationDispatch: dict[ActivationType, type[Activation]] = { - "clamped_swiglu": ClampedSwiglu, "silu": Silu, "relu2": Relu2, } diff --git a/src/prime_rl/trainer/models/layers/attn.py b/src/prime_rl/trainer/models/layers/attn.py index ef9bf59a0c1..37ce8d9a385 100644 --- a/src/prime_rl/trainer/models/layers/attn.py +++ b/src/prime_rl/trainer/models/layers/attn.py @@ -307,7 +307,3 @@ def _ring_compute_attention(self, q, k, v, cu_seqlens, max_seqlen): from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention AfmoeFlashAttention._compute_attention = _ring_compute_attention - - from prime_rl.trainer.models.gpt_oss.attention import substitute_gpt_oss_ring_attention - - substitute_gpt_oss_ring_attention(process_group, heads_k_stride) diff --git a/src/prime_rl/trainer/models/layers/expert_compute.py b/src/prime_rl/trainer/models/layers/expert_compute.py index 700667a3593..17a255c7d91 100644 --- a/src/prime_rl/trainer/models/layers/expert_compute.py +++ b/src/prime_rl/trainer/models/layers/expert_compute.py @@ -26,21 +26,6 @@ def __call__( ) -> torch.Tensor: ... -def broadcast_expert_bias( - bias: torch.Tensor, - num_tokens_per_expert: torch.Tensor, - target_rows: int, -) -> torch.Tensor: - repeats = num_tokens_per_expert.to(torch.int64) - padding_rows = repeats.new_tensor(target_rows) - repeats.sum() - return torch.repeat_interleave( - torch.cat((bias, bias.new_zeros((1, bias.shape[1])))), - torch.cat((repeats, padding_rows.unsqueeze(0))), - dim=0, - output_size=target_rows, - ) - - class GroupedGemmExpertCompute: def __init__( self, @@ -51,7 +36,7 @@ def __init__( self.token_group_alignment = token_group_alignment def validate(self, experts: "GroupedExperts") -> None: - """The shared forward handles the experts' activation, biases, and weight layout.""" + """The shared forward handles the experts' activation and weight layout.""" def __call__(self, experts: "GroupedExperts", x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: assert x.dim() == 2 @@ -75,20 +60,9 @@ def to_local(tensor: torch.Tensor) -> torch.Tensor: gate_up = self.gemm(x_bf16, gate_up_proj.bfloat16(), offsets) gate, up = gate_up.chunk(2, dim=-1) - if experts.up_proj_bias is not None: - up_proj_bias = to_local(experts.up_proj_bias) - up = up + broadcast_expert_bias(up_proj_bias, num_tokens_per_expert, up.shape[0]).bfloat16() - - if gate is not None and experts.gate_proj_bias is not None: - gate_proj_bias = to_local(experts.gate_proj_bias) - gate = gate + broadcast_expert_bias(gate_proj_bias, num_tokens_per_expert, gate.shape[0]).bfloat16() - hidden = experts.activation.apply(gate, up) down_proj = to_local(experts.down_proj).transpose(-2, -1) output = self.gemm(hidden, down_proj.bfloat16(), offsets) - if experts.down_proj_bias is not None: - down_proj_bias = to_local(experts.down_proj_bias) - output = output + broadcast_expert_bias(down_proj_bias, num_tokens_per_expert, output.shape[0]).bfloat16() return output.type_as(x) diff --git a/src/prime_rl/trainer/models/layers/lora/multi_moe.py b/src/prime_rl/trainer/models/layers/lora/multi_moe.py index c87201ddef4..2227f8033d9 100644 --- a/src/prime_rl/trainer/models/layers/lora/multi_moe.py +++ b/src/prime_rl/trainer/models/layers/lora/multi_moe.py @@ -4,8 +4,6 @@ from torch import nn from torch.distributed.tensor import DTensor -from prime_rl.trainer.models.layers.activations import ClampedSwiglu -from prime_rl.trainer.models.layers.expert_compute import broadcast_expert_bias from prime_rl.trainer.models.layers.lora.base import MultiLoRAModule, get_lora_num_tokens, get_multilora_scaling from prime_rl.trainer.models.layers.moe import GroupedExperts @@ -583,257 +581,3 @@ def __repr__(self) -> str: f"n_adapters={self.n_adapters}, num_experts={self.num_experts}, " f"alpha={self.alpha}, dropout={self.lora_dropout})" ) - - -class MultiLoRAGptOssGroupedExperts(MultiLoRAModule): - """ - GPT-OSS GroupedExperts + multi-LoRA. - - Preserves GPT-OSS's combined gate/up adapter format while applying it to the - canonical split gate_proj/up_proj runtime weights. - """ - - def __init__( - self, - base_layer: GroupedExperts, - rank: int, - n_adapters: int, - alpha: float = 32.0, - dropout: float = 0.0, - ): - super().__init__(base_layer) - if rank <= 0 or n_adapters <= 0: - raise ValueError("rank and n_adapters must be > 0") - if base_layer.gate_proj is None or base_layer.activation is not ClampedSwiglu: - raise ValueError("MultiLoRAGptOssGroupedExperts requires gated GPT-OSS experts") - if any( - bias is None for bias in (base_layer.gate_proj_bias, base_layer.up_proj_bias, base_layer.down_proj_bias) - ): - raise ValueError("GPT-OSS experts require projection biases") - - self.num_experts = base_layer.num_experts - self.hidden_size = base_layer.up_proj.shape[2] - self.intermediate_size = base_layer.up_proj.shape[1] - self.gate_up_out = 2 * self.intermediate_size - - if rank % 8 != 0 or self.hidden_size % 8 != 0 or self.intermediate_size % 8 != 0: - raise ValueError("grouped_mm requires rank and expert dimensions divisible by 8") - - self.rank = rank - self.n_adapters = n_adapters - self.alpha = alpha - self.lora_dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity() - - self._lora_num_tokens = get_lora_num_tokens() - self._scaling_factors = get_multilora_scaling() - - self.gate_up_lora_A = nn.ParameterList( - [ - nn.Parameter( - torch.empty( - self.num_experts, - rank, - self.hidden_size, - device=base_layer.up_proj.device, - dtype=base_layer.up_proj.dtype, - ) - ) - for _ in range(n_adapters) - ] - ) - self.gate_up_lora_B = nn.ParameterList( - [ - nn.Parameter( - torch.empty( - self.num_experts, - self.gate_up_out, - rank, - device=base_layer.up_proj.device, - dtype=base_layer.up_proj.dtype, - ) - ) - for _ in range(n_adapters) - ] - ) - - self.down_lora_A = nn.ParameterList( - [ - nn.Parameter( - torch.empty( - self.num_experts, - rank, - self.intermediate_size, - device=base_layer.down_proj.device, - dtype=base_layer.down_proj.dtype, - ) - ) - for _ in range(n_adapters) - ] - ) - self.down_lora_B = nn.ParameterList( - [ - nn.Parameter( - torch.empty( - self.num_experts, - self.hidden_size, - rank, - device=base_layer.down_proj.device, - dtype=base_layer.down_proj.dtype, - ) - ) - for _ in range(n_adapters) - ] - ) - - self.reset_parameters() - - def reset_parameters(self, index: int | None = None) -> None: - if index is None: - for i in range(self.n_adapters): - self.reset_parameters(i) - else: - nn.init.kaiming_uniform_(self.gate_up_lora_A[index], a=math.sqrt(5)) - nn.init.zeros_(self.gate_up_lora_B[index]) - nn.init.kaiming_uniform_(self.down_lora_A[index], a=math.sqrt(5)) - nn.init.zeros_(self.down_lora_B[index]) - - def named_parameters_for_adapter(self, idx: int) -> list[tuple[str, nn.Parameter]]: - return [ - ("gate_up_lora_A", self.gate_up_lora_A[idx]), - ("gate_up_lora_B", self.gate_up_lora_B[idx]), - ("down_lora_A", self.down_lora_A[idx]), - ("down_lora_B", self.down_lora_B[idx]), - ] - - def get_lora_param_counts(self) -> tuple[int, int]: - adapter_params = ( - self.gate_up_lora_A[0].numel() - + self.gate_up_lora_B[0].numel() - + self.down_lora_A[0].numel() - + self.down_lora_B[0].numel() - ) - adapted_params = ( - self.base_layer.gate_proj.numel() + self.base_layer.up_proj.numel() + self.base_layer.down_proj.numel() - ) - return adapter_params, adapted_params - - def state_dict_for_adapter(self, idx: int) -> dict[str, torch.Tensor]: - """vLLM-compatible 3D MoE adapter format. - - For 3D MoE models (gpt-oss in vLLM has `is_3d_moe_weight = True`), vLLM expects: - - `experts.base_layer.lora_{A,B}.weight` for the gate_up projection - - `experts.lora_{A,B}.weight` for the down projection - with experts stacked into the rank dim. See - vllm/lora/model_manager.py::_stack_moe_lora_weights, which reshapes - lora_A: (num_experts*rank, in) -> (num_experts, rank, in) - lora_B: (out, rank*num_experts) -> (out, rank, num_experts) -> (num_experts, out, rank) - """ - detached_gu_a = self.gate_up_lora_A[idx].detach() - detached_gu_b = self.gate_up_lora_B[idx].detach() - detached_d_a = self.down_lora_A[idx].detach() - detached_d_b = self.down_lora_B[idx].detach() - - if isinstance(detached_gu_a, DTensor): - detached_gu_a = detached_gu_a.full_tensor() - detached_gu_b = detached_gu_b.full_tensor() - detached_d_a = detached_d_a.full_tensor() - detached_d_b = detached_d_b.full_tensor() - - # lora_A: (num_experts, rank, in) -> (num_experts*rank, in) - gu_a_flat = detached_gu_a.reshape(self.num_experts * self.rank, self.hidden_size).clone() - d_a_flat = detached_d_a.reshape(self.num_experts * self.rank, self.intermediate_size).clone() - # lora_B: (num_experts, out, rank) -> (out, rank, num_experts) -> (out, rank*num_experts) - # vLLM's reshape treats the last dim of lora_B as (rank, num_experts) with experts fast-varying. - gu_b_flat = detached_gu_b.permute(1, 2, 0).contiguous().reshape(self.gate_up_out, self.rank * self.num_experts) - d_b_flat = detached_d_b.permute(1, 2, 0).contiguous().reshape(self.hidden_size, self.rank * self.num_experts) - - return { - "base_layer.lora_A.weight": gu_a_flat, - "base_layer.lora_B.weight": gu_b_flat, - "lora_A.weight": d_a_flat, - "lora_B.weight": d_b_flat, - } - - def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: - adapter_idx = self._lora_num_tokens.argmax().item() - gu_a = self.gate_up_lora_A[adapter_idx] - gu_b = self.gate_up_lora_B[adapter_idx] - d_a = self.down_lora_A[adapter_idx] - d_b = self.down_lora_B[adapter_idx] - - scaling = self._scaling_factors[adapter_idx].item() - - gate_proj = self.base_layer.gate_proj - up_proj = self.base_layer.up_proj - down_proj = self.base_layer.down_proj - gate_proj_bias = self.base_layer.gate_proj_bias - up_proj_bias = self.base_layer.up_proj_bias - down_proj_bias = self.base_layer.down_proj_bias - - permuted_indices = None - if isinstance(up_proj, DTensor): - gate_proj = gate_proj.to_local() - up_proj = up_proj.to_local() - down_proj = down_proj.to_local() - gate_proj_bias = gate_proj_bias.to_local() - up_proj_bias = up_proj_bias.to_local() - down_proj_bias = down_proj_bias.to_local() - gu_a = gu_a.to_local() - gu_b = gu_b.to_local() - d_a = d_a.to_local() - d_b = d_b.to_local() - - if getattr(self.base_layer, "ep_comm_backend", "torch") != "deepep": - from torchtitan.experiments.kernels.moe.indices import generate_permute_indices - - from prime_rl.trainer.distributed.expert_parallel import TOKEN_GROUP_ALIGN_SIZE_M - - experts_per_ep_rank = up_proj.shape[0] - num_ep_ranks = num_tokens_per_expert.shape[0] // experts_per_ep_rank - - with torch.no_grad(): - permuted_indices, num_tokens_per_expert, _ = generate_permute_indices( - num_tokens_per_expert, - experts_per_ep_rank, - num_ep_ranks, - x.shape[0] + experts_per_ep_rank * TOKEN_GROUP_ALIGN_SIZE_M, - TOKEN_GROUP_ALIGN_SIZE_M, - ) - - x = torch.vstack((x, x.new_zeros((x.shape[-1])))) - input_shape = x.shape - x = x[permuted_indices, :] - - offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) - lora_x = self.lora_dropout(x) - - gate = torch._grouped_mm(x.bfloat16(), gate_proj.bfloat16().transpose(-2, -1), offs=offsets) - gate = gate + broadcast_expert_bias(gate_proj_bias, num_tokens_per_expert, gate.shape[0]).bfloat16() - up = torch._grouped_mm(x.bfloat16(), up_proj.bfloat16().transpose(-2, -1), offs=offsets) - up = up + broadcast_expert_bias(up_proj_bias, num_tokens_per_expert, up.shape[0]).bfloat16() - - gate_up_lora = _run_lora_grouped_mm(lora_x, gu_a, gu_b, offsets) - gate = gate + scaling * gate_up_lora[..., ::2].bfloat16() - up = up + scaling * gate_up_lora[..., 1::2].bfloat16() - - h = self.base_layer.activation.apply(gate, up) - lora_h = self.lora_dropout(h) - - out_base = torch._grouped_mm(h, down_proj.bfloat16().transpose(-2, -1), offs=offsets) - out_base = out_base + broadcast_expert_bias(down_proj_bias, num_tokens_per_expert, out_base.shape[0]).bfloat16() - out_lora = _run_lora_grouped_mm(lora_h, d_a, d_b, offsets) - out = (out_base + scaling * out_lora.bfloat16()).type_as(x) - - if permuted_indices is not None: - out_unpermuted = out.new_zeros(input_shape) - out_unpermuted[permuted_indices, :] = out - out = out_unpermuted[:-1] - - return out - - def __repr__(self) -> str: - return ( - f"{self.__class__.__name__}(base={self.base_layer}, rank={self.rank}, " - f"n_adapters={self.n_adapters}, num_experts={self.num_experts}, " - f"alpha={self.alpha}, dropout={self.lora_dropout})" - ) diff --git a/src/prime_rl/trainer/models/layers/moe.py b/src/prime_rl/trainer/models/layers/moe.py index 99a25cb4e4a..cee40d725c6 100644 --- a/src/prime_rl/trainer/models/layers/moe.py +++ b/src/prime_rl/trainer/models/layers/moe.py @@ -17,7 +17,7 @@ from prime_rl.trainer.models.layers.expert_compute import BF16ExpertCompute, ExpertCompute from prime_rl.trainer.models.layers.mlp import ExpertType, FeedForward -ScoreFuncType = Literal["softmax", "sigmoid", "topk_softmax"] +ScoreFuncType = Literal["softmax", "sigmoid"] @torch.library.custom_op( @@ -77,7 +77,6 @@ def __init__( *, expert_type: ExpertType = "gated", activation: ActivationType = "silu", - bias: bool = False, compute: ExpertCompute | None = None, ): super().__init__() @@ -87,11 +86,6 @@ def __init__( self.up_proj = nn.Parameter(torch.empty(num_experts, hidden_dim, dim)) self.register_parameter("gate_up_proj", None) self.down_proj = nn.Parameter(torch.empty(num_experts, dim, hidden_dim)) - self.gate_proj_bias = ( - nn.Parameter(torch.empty(num_experts, hidden_dim)) if bias and self.gate_proj is not None else None - ) - self.up_proj_bias = nn.Parameter(torch.empty(num_experts, hidden_dim)) if bias else None - self.down_proj_bias = nn.Parameter(torch.empty(num_experts, dim)) if bias else None self.activation = ActivationDispatch[activation] if expert_type == "non_gated": @@ -124,9 +118,6 @@ def init_weights(self, init_std: float): remaining = (up_proj, self.down_proj) for weight in remaining: nn.init.trunc_normal_(weight, mean=0.0, std=init_std) - for bias in (self.gate_proj_bias, self.up_proj_bias, self.down_proj_bias): - if bias is not None: - nn.init.zeros_(bias) class TokenChoiceTopKRouter(nn.Module): @@ -136,11 +127,9 @@ class TokenChoiceTopKRouter(nn.Module): dim (int): Dimension of input tokens. num_experts (int): Number of experts in each moe layer. top_k (int): Number of experts each token will be routed to in token-choice routing. - score_func (Literal["softmax", "sigmoid", "topk_softmax"]): Score transform. ``topk_softmax`` - selects experts from the logits and normalizes only the selected logits. + score_func (Literal["softmax", "sigmoid"]): Score transform. route_norm (bool): Whether to normalize the routing scores when using sigmoid. route_scale (float): Scaling factor applied to the routing scores. - gate_bias (bool): Whether the gate has a trainable logit bias. selection_bias (bool): Whether to keep a persistent selection-only bias. The bias affects expert selection but not routing weights. topk_sorted (bool): Whether selected experts are returned in descending score order. @@ -151,16 +140,15 @@ def __init__( dim: int, num_experts: int, top_k: int, - score_func: Literal["softmax", "sigmoid", "topk_softmax"], + score_func: ScoreFuncType, route_norm: bool, route_scale: float, *, - gate_bias: bool = False, selection_bias: bool = False, topk_sorted: bool = True, ): super().__init__() - self.gate = nn.Linear(dim, num_experts, bias=gate_bias) + self.gate = nn.Linear(dim, num_experts, bias=False) self.register_buffer( "selection_bias", torch.zeros(num_experts, dtype=torch.float32) if selection_bias else None, @@ -202,8 +190,7 @@ def forward( f"routed_experts shape: {routed_experts.shape}, top_k: {self.top_k}" ) if self.fp32_gate: - gate_bias = self.gate.bias.float() if self.gate.bias is not None else None - logits = F.linear(x.float(), self.gate.weight.float(), gate_bias) + logits = F.linear(x.float(), self.gate.weight.float()) else: logits = self.gate(x) @@ -212,8 +199,6 @@ def forward( scores = torch.sigmoid(logits.float()) elif self.score_func == "softmax": scores = F.softmax(logits.float(), dim=1) - elif self.score_func == "topk_softmax": - scores = logits else: raise NotImplementedError(f"Unknown score function {self.score_func}") @@ -241,11 +226,8 @@ def forward( ) top_scores = scores.gather(dim=1, index=selected_experts_indices) - if self.score_func == "topk_softmax": - top_scores = F.softmax(top_scores, dim=-1, dtype=top_scores.dtype) - with torch.no_grad(): - if self.score_func in ("softmax", "topk_softmax"): + if self.score_func == "softmax": routing_confidence_sum = top_scores.sum() else: selected_probability_mass = top_scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) @@ -268,8 +250,6 @@ def forward( def init_weights(self, init_std: float): nn.init.trunc_normal_(self.gate.weight, mean=0.0, std=init_std) - if self.gate.bias is not None: - nn.init.zeros_(self.gate.bias) class MoE(nn.Module): diff --git a/src/prime_rl/trainer/models/layers/sonic_moe.py b/src/prime_rl/trainer/models/layers/sonic_moe.py index f9e0e95e280..25d5b6aaf09 100644 --- a/src/prime_rl/trainer/models/layers/sonic_moe.py +++ b/src/prime_rl/trainer/models/layers/sonic_moe.py @@ -20,11 +20,7 @@ def validate(self, experts: "GroupedExperts") -> None: "SonicMoE requires gated experts with fused gate/up projections. Enable the gate_up fusion." ) if experts.activation is not Silu: - raise ValueError( - "SonicMoE requires standard SwiGLU (silu); custom or clamped activations are not supported." - ) - if any(bias is not None for bias in (experts.gate_proj_bias, experts.up_proj_bias, experts.down_proj_bias)): - raise ValueError("SonicMoE requires bias-free experts.") + raise ValueError("SonicMoE requires standard SwiGLU (silu).") def __call__(self, experts: "GroupedExperts", x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: gate_up = experts.gate_up_proj diff --git a/src/prime_rl/trainer/models/layers/ulysses_attn.py b/src/prime_rl/trainer/models/layers/ulysses_attn.py index 9f2c855e294..d93c292c1f4 100644 --- a/src/prime_rl/trainer/models/layers/ulysses_attn.py +++ b/src/prime_rl/trainer/models/layers/ulysses_attn.py @@ -116,10 +116,8 @@ def ulysses_flash_attn_varlen_func( cp_size: int, flash_attn_version: int = 2, window_size: tuple[int, int] = (-1, -1), - softmax_scale: float | None = None, dropout_p: float = 0.0, deterministic: bool | None = None, - learnable_sink: torch.Tensor | None = None, ) -> torch.Tensor: """Run varlen flash attention under Ulysses CP. @@ -140,16 +138,10 @@ def ulysses_flash_attn_varlen_func( kwargs: dict = {"causal": causal} if window_size != (-1, -1): kwargs["window_size"] = window_size - if softmax_scale is not None: - kwargs["softmax_scale"] = softmax_scale if dropout_p: kwargs["dropout_p"] = dropout_p if deterministic is not None: kwargs["deterministic"] = deterministic - if learnable_sink is not None: - local_heads = q.shape[1] - rank = dist.get_rank(cp_group) - kwargs["learnable_sink"] = learnable_sink[rank * local_heads : (rank + 1) * local_heads] if flash_attn_version == 4: # FA4 takes cu_seqlens as keyword args (qv positional collides otherwise). @@ -222,7 +214,3 @@ def _ulysses_compute_attention(self, q, k, v, cu_seqlens, max_seqlen): from prime_rl.trainer.models.afmoe.modeling_afmoe import AfmoeFlashAttention AfmoeFlashAttention._compute_attention = _ulysses_compute_attention - - from prime_rl.trainer.models.gpt_oss.attention import substitute_gpt_oss_ulysses_attention - - substitute_gpt_oss_ulysses_attention(process_group) diff --git a/tests/unit/test_parsers.py b/tests/unit/test_parsers.py index 584fcb5e53d..351f3c1ab25 100644 --- a/tests/unit/test_parsers.py +++ b/tests/unit/test_parsers.py @@ -41,9 +41,6 @@ ("MiniMaxAI/MiniMax-M2", "minimax_m2", "minimax_m2_append_think"), ("MiniMaxAI/MiniMax-M2.1", "minimax_m2", "minimax_m2_append_think"), ("MiniMaxAI/MiniMax-M2.5", "minimax_m2", "minimax_m2_append_think"), - # gpt-oss (reasoning handled natively by vLLM's harmony path) - ("openai/gpt-oss-20b", "openai", None), - ("openai/gpt-oss-120b", "openai", None), # Poolside Laguna ("poolside/Laguna-S-2.1", "poolside_v1", "poolside_v1"), ("poolside/Laguna-S-2.1-FP8", "poolside_v1", "poolside_v1"), diff --git a/tests/unit/train/models/test_gpt_oss.py b/tests/unit/train/models/test_gpt_oss.py deleted file mode 100644 index 2fc4ec9ee70..00000000000 --- a/tests/unit/train/models/test_gpt_oss.py +++ /dev/null @@ -1,133 +0,0 @@ -import os - -import pytest -import torch -import torch.distributed as dist - -from prime_rl.trainer.models.gpt_oss import GptOssConfig -from prime_rl.trainer.models.gpt_oss import GptOssForCausalLM as PrimeRLGptOssForCausalLM -from prime_rl.trainer.models.gpt_oss.attention import ( - GptOssAttention, - substitute_gpt_oss_ring_attention, - substitute_gpt_oss_ulysses_attention, -) -from prime_rl.utils.cp import setup_cp_attention_params - - -def _config() -> GptOssConfig: - return GptOssConfig( - num_hidden_layers=1, - num_local_experts=4, - vocab_size=128, - hidden_size=64, - intermediate_size=32, - head_dim=16, - num_attention_heads=4, - num_key_value_heads=2, - max_position_embeddings=128, - num_experts_per_tok=2, - sliding_window=4, - rope_parameters={"rope_type": "default", "rope_theta": 150000.0}, - attn_implementation="flash_attention_4", - use_cache=False, - ) - - -def test_gpt_oss_checkpoint_conversion_roundtrip(): - with torch.device("meta"): - model = PrimeRLGptOssForCausalLM(_config()) - - prime_state_dict = {name: torch.randn(tensor.shape) for name, tensor in model.state_dict().items()} - expected_prime = {name: tensor.clone() for name, tensor in prime_state_dict.items()} - hf_state_dict = model.convert_to_hf(prime_state_dict) - - assert model.is_hf_state_dict(hf_state_dict) - assert not model.is_prime_state_dict(hf_state_dict) - - roundtrip = model.convert_to_prime(hf_state_dict) - assert roundtrip.keys() == expected_prime.keys() - for name, tensor in roundtrip.items(): - torch.testing.assert_close(tensor, expected_prime[name]) - - -@pytest.fixture(scope="module") -def cp_process_group(): - if int(os.environ.get("WORLD_SIZE", 1)) != 2: - pytest.skip("run with torchrun --nproc-per-node=2") - if torch.cuda.get_device_capability()[0] not in (9, 10, 11): - pytest.skip("GPT-OSS learned sinks require SM90 or SM100/SM110") - - local_rank = int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) - dist.init_process_group("nccl") - yield dist.group.WORLD - dist.destroy_process_group() - - -@pytest.mark.gpu -@pytest.mark.parametrize("cp_style", ["ring", "ulysses"]) -def test_gpt_oss_context_parallel_attention(cp_style: str, cp_process_group): - local_rank = int(os.environ["LOCAL_RANK"]) - process_group = cp_process_group - original_compute_attention = GptOssAttention.compute_attention - - try: - torch.manual_seed(0) - config = _config() - config.head_dim = 64 - attention = GptOssAttention(config, layer_idx=0).cuda().to(torch.bfloat16) - query = torch.randn(8, 4, 64, device="cuda", dtype=torch.bfloat16) - key = torch.randn(8, 2, 64, device="cuda", dtype=torch.bfloat16) - value = torch.randn(8, 2, 64, device="cuda", dtype=torch.bfloat16) - with torch.no_grad(): - for tensor in (attention.sinks, query, key, value): - dist.broadcast(tensor, src=0) - cu_seqlens = torch.tensor([0, 5, 8], device="cuda", dtype=torch.int32) - - reference_query = query.detach().clone().requires_grad_() - reference_key = key.detach().clone().requires_grad_() - reference_value = value.detach().clone().requires_grad_() - reference = original_compute_attention( - attention, - reference_query, - reference_key, - reference_value, - cu_seqlens, - 5, - ) - reference_output = reference.detach().clone() - reference.float().square().sum().backward() - reference_sink_grad = attention.sinks.grad.detach().clone() - attention.sinks.grad = None - - position_ids = torch.tensor([[0, 1, 2, 3, 4, 0, 1, 2]], device="cuda") - setup_cp_attention_params( - position_ids, - process_group, - seq_lens=torch.tensor([5, 3], device="cuda"), - cp_style=cp_style, - ) - if cp_style == "ring": - substitute_gpt_oss_ring_attention(process_group, heads_k_stride=1) - else: - substitute_gpt_oss_ulysses_attention(process_group) - - local_slice = slice(local_rank * 4, (local_rank + 1) * 4) - local_query = query[local_slice].detach().clone().requires_grad_() - local_key = key[local_slice].detach().clone().requires_grad_() - local_value = value[local_slice].detach().clone().requires_grad_() - actual = attention.compute_attention(local_query, local_key, local_value, cu_seqlens, 5) - actual.float().square().sum().backward() - - gathered = [torch.empty_like(actual) for _ in range(2)] - dist.all_gather(gathered, actual) - torch.testing.assert_close(torch.cat(gathered), reference_output, rtol=3e-2, atol=3e-2) - torch.testing.assert_close(local_query.grad, reference_query.grad[local_slice], rtol=5e-2, atol=5e-2) - torch.testing.assert_close(local_key.grad, reference_key.grad[local_slice], rtol=5e-2, atol=5e-2) - torch.testing.assert_close(local_value.grad, reference_value.grad[local_slice], rtol=5e-2, atol=5e-2) - - sink_grad = attention.sinks.grad - dist.all_reduce(sink_grad) - torch.testing.assert_close(sink_grad, reference_sink_grad, rtol=5e-2, atol=5e-2) - finally: - GptOssAttention.compute_attention = original_compute_attention diff --git a/tests/unit/train/models/test_moe.py b/tests/unit/train/models/test_moe.py index 86d477b7ace..cae926d895b 100644 --- a/tests/unit/train/models/test_moe.py +++ b/tests/unit/train/models/test_moe.py @@ -67,8 +67,6 @@ def test_unselected_moe_uses_bf16_without_loading_compute_backend(selection, com pytest.param(GroupedExperts, {}, False, "fused gate/up", id="unfused"), pytest.param(GroupedExperts, {"expert_type": "non_gated"}, False, "gated experts", id="non-gated"), pytest.param(GroupedExperts, {"activation": "relu2"}, True, "standard SwiGLU", id="relu2"), - pytest.param(GroupedExperts, {"activation": "clamped_swiglu"}, True, "standard SwiGLU", id="clamped-swiglu"), - pytest.param(GroupedExperts, {"bias": True}, True, "bias-free", id="bias"), pytest.param(DeepseekV4Experts, {"swiglu_limit": 10.0}, True, "standard SwiGLU", id="deepseek-v4"), ], ) @@ -175,7 +173,6 @@ def test_expert_type_and_activation_are_independent(expert_type, activation): num_experts=2, expert_type=expert_type, activation=activation, - bias=True, ) experts.set_compute(GroupedGemmExpertCompute(_grouped_mm_reference, token_group_alignment=1)) assert experts.token_group_alignment == 1 @@ -183,7 +180,6 @@ def test_expert_type_and_activation_are_independent(expert_type, activation): has_gate = expert_type == "gated" assert (experts.gate_proj is not None) == has_gate - assert (experts.gate_proj_bias is not None) == has_gate assert ("gate_proj" in experts.state_dict()) == has_gate x = torch.randn(3, 4) @@ -194,20 +190,14 @@ def test_expert_type_and_activation_are_independent(expert_type, activation): start = 0 for expert, count in enumerate(counts.tolist()): expert_input = x[start : start + count].bfloat16() - up = F.linear(expert_input, experts.up_proj[expert].bfloat16(), experts.up_proj_bias[expert].bfloat16()) + up = F.linear(expert_input, experts.up_proj[expert].bfloat16()) activation_input = up if has_gate: - activation_input = F.linear( - expert_input, - experts.gate_proj[expert].bfloat16(), - experts.gate_proj_bias[expert].bfloat16(), - ) + activation_input = F.linear(expert_input, experts.gate_proj[expert].bfloat16()) hidden = F.silu(activation_input) if activation == "silu" else F.relu(activation_input).square() if has_gate: hidden = hidden * up - expected.append( - F.linear(hidden, experts.down_proj[expert].bfloat16(), experts.down_proj_bias[expert].bfloat16()) - ) + expected.append(F.linear(hidden, experts.down_proj[expert].bfloat16())) start += count torch.testing.assert_close(actual, torch.cat(expected).float())