diff --git a/core/parallel/expert_module.cpp b/core/parallel/expert_module.cpp index dd336985..c92a1c9e 100644 --- a/core/parallel/expert_module.cpp +++ b/core/parallel/expert_module.cpp @@ -196,6 +196,8 @@ torch::Tensor MoEMLP::forward(torch::Tensor hidden_states, DLOG_FATAL_IF(batch_size > kMaxTokens || batch_size <= 0, "batch_size should be (0,", kMaxTokens, "] , but got", batch_size); + TORCH_CHECK(hidden_states.scalar_type() == input_.scalar_type(), + "hidden_states dtype must match expert input dtype"); // Use async copy with the provided execution stream cudaMemcpyAsync(input_.data_ptr(), hidden_states.data_ptr(), diff --git a/moe_infinity/models/gpt_oss.py b/moe_infinity/models/gpt_oss.py index 2f8cbdee..9bec1388 100644 --- a/moe_infinity/models/gpt_oss.py +++ b/moe_infinity/models/gpt_oss.py @@ -36,6 +36,13 @@ def __init__(self, config): class SyncGptOssMLP(nn.Module): archer_config: Optional[ArcherConfig] = None layer_id: Optional[int] = None + expert_executor = None + expert_prefetcher = None + expert_tracer = None + expert_predictor = None + expert_tensor_map = None + archer_engine = None + lib = None def __init__(self, config): super().__init__() diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index d1a7c393..569ae5e5 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -174,6 +174,33 @@ def _out_block(experts_prefix: str) -> str: "down_proj_bias", ) +_GPT_OSS_RESIDENT_RATIO = 0.9 + + +def _gpt_oss_offload_enabled(model_type, archer_config): + if model_type != "gpt_oss": + return True + return archer_config.device_memory_ratio < _GPT_OSS_RESIDENT_RATIO + + +def _wire_expert_collaborators( + module, + enabled, + *, + expert_executor, + expert_prefetcher, + expert_tracer, + expert_predictor, + expert_tensor_map, +): + if not enabled: + return + module.expert_executor = expert_executor + module.expert_prefetcher = expert_prefetcher + module.expert_tracer = expert_tracer + module.expert_predictor = expert_predictor + module.expert_tensor_map = expert_tensor_map + def _expand_gpt_oss_packed_experts(state_dict, config): if getattr(config, "model_type", "") != "gpt_oss": @@ -542,6 +569,9 @@ def init( ) self.archer_config = _archer_config + self.gpt_oss_offload_enabled = _gpt_oss_offload_enabled( + getattr(self.config, "model_type", ""), self.archer_config + ) if _archer_config.trace_path is not None: self.expert_tracer.load_trace(_archer_config.trace_path) @@ -1078,12 +1108,19 @@ def archer_from_pretrained(cls, *args, **kwargs): module.is_gptq = is_gptq_quantized(self.config) self.expert_modules.append(module) - if not isinstance(module, SyncGptOssMLP): - module.expert_executor = self.expert_executor - module.expert_prefetcher = self.expert_prefetcher - module.expert_tracer = self.expert_tracer - module.expert_predictor = self.expert_predictor - module.expert_tensor_map = self.expert_tensor_map + use_dispatcher = ( + not isinstance(module, SyncGptOssMLP) + or self.gpt_oss_offload_enabled + ) + _wire_expert_collaborators( + module, + use_dispatcher, + expert_executor=self.expert_executor, + expert_prefetcher=self.expert_prefetcher, + expert_tracer=self.expert_tracer, + expert_predictor=self.expert_predictor, + expert_tensor_map=self.expert_tensor_map, + ) module.lib = self.prefetch_lib @@ -1101,7 +1138,10 @@ def archer_from_pretrained(cls, *args, **kwargs): if self._is_shared_expert_param(_name): del self.name_id_map[_name] - if getattr(self.config, "model_type", "") == "gpt_oss": + if ( + getattr(self.config, "model_type", "") == "gpt_oss" + and not self.gpt_oss_offload_enabled + ): self._load_resident_gpt_oss(model) self.setup_archer_hooks(model) diff --git a/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py index bbb13cfc..c5ec2d0c 100644 --- a/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py +++ b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py @@ -129,3 +129,137 @@ def test_dequantized_option_a_matches_resident_expert_forward(): ) assert ((option_a.float() - golden).abs() <= envelope).all() assert ((resident.float() - golden).abs() <= envelope).all() + + +def _native_dispatch(tmp_path, hidden_dtype): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + try: + from moe_infinity.runtime.model_offload import _load_prefetch_lib + + prefetch_lib = _load_prefetch_lib() + except Exception as exc: + pytest.skip(f"native Archer extension unavailable: {exc}") + + from moe_infinity.models.gpt_oss import SyncGptOssMLP + + class Config: + hidden_size = 64 + intermediate_size = 32 + num_local_experts = 2 + num_experts_per_tok = 1 + + torch.manual_seed(137) + hidden = Config.hidden_size + intermediate = Config.intermediate_size + mlp = SyncGptOssMLP(Config()) + tensors = [ + torch.randint( + 0, 256, (2 * intermediate, hidden // 2), dtype=torch.uint8 + ), + torch.randint( + 120, 135, (2 * intermediate, hidden // 32), dtype=torch.uint8 + ), + torch.randn(2 * intermediate, dtype=torch.bfloat16), + torch.randint(0, 256, (hidden, intermediate // 2), dtype=torch.uint8), + torch.randint( + 120, 135, (hidden, intermediate // 32), dtype=torch.uint8 + ), + torch.randn(hidden, dtype=torch.bfloat16), + ] + mlp.experts.gate_up_proj.requires_grad_(False) + mlp.experts.down_proj.requires_grad_(False) + mlp.experts.gate_up_proj.data = tensors[0].unsqueeze(0).repeat(2, 1, 1) + mlp.experts.gate_up_proj_scales.data = ( + tensors[1].unsqueeze(0).repeat(2, 1, 1) + ) + mlp.experts.gate_up_proj_bias.data = tensors[2].unsqueeze(0).repeat(2, 1) + mlp.experts.down_proj.data = tensors[3].unsqueeze(0).repeat(2, 1, 1) + mlp.experts.down_proj_scales.data = tensors[4].unsqueeze(0).repeat(2, 1, 1) + mlp.experts.down_proj_bias.data = tensors[5].unsqueeze(0).repeat(2, 1) + + hidden_states = torch.randn(3, hidden, dtype=hidden_dtype, device="cuda:0") + expected = mlp._expert_forward_mxfp4(hidden_states, 0) + + engine = prefetch_lib.prefetch_handle(f"{tmp_path}/", 0.5) + expert_tensor_ids = [list(range(6)), list(range(6, 12))] + for tensor_id, tensor in zip(expert_tensor_ids[0], tensors): + engine.offload(tensor, tensor_id) + for tensor_id, tensor in zip(expert_tensor_ids[1], tensors): + engine.offload(tensor, tensor_id) + dense_tensor_ids = list(range(12, 12 + 2 * torch.cuda.device_count())) + for tensor_id in dense_tensor_ids: + engine.offload(torch.zeros(1, dtype=torch.bfloat16), tensor_id) + split = len(dense_tensor_ids) // 2 + dense_before = [ + (f"model.dense.{index}", [[tensor_id]]) + for index, tensor_id in enumerate(dense_tensor_ids[:split]) + ] + dense_after = [ + (f"model.dense.{index + split}", [[tensor_id]]) + for index, tensor_id in enumerate(dense_tensor_ids[split:]) + ] + engine.set_topology( + dense_before + + [("model.layers.0.mlp.experts", expert_tensor_ids)] + + dense_after + ) + + dispatcher = prefetch_lib.expert_dispatcher(2, 1, 0, 6, 1) + dispatcher.register_expert(0, 0, expert_tensor_ids[0], "") + dispatcher.register_expert(0, 1, expert_tensor_ids[1], "") + torch.cuda.set_device(hidden_states.device) + router_mask = torch.zeros((3, 2), dtype=torch.bool, device="cuda:0") + router_mask[:, 0] = True + router_weights = torch.zeros((3, 2), dtype=torch.bfloat16, device="cuda:0") + router_weights[:, 0] = 1 + dispatcher.set_inputs( + hidden_states, + router_mask, + router_weights, + ) + dispatcher.set_expected_queue(1) + dispatcher.enqueue_expert(0, 0, 0, False) + dispatcher.notify_fetch_start() + actual = dispatcher.wait_expert() + + return actual, expected, hidden_states, tensors + + +@pytest.mark.gpu +def test_native_dispatcher_matches_resident_expert_forward(tmp_path): + actual, resident, hidden_states, tensors = _native_dispatch( + tmp_path, torch.bfloat16 + ) + from moe_infinity._v4_fp4 import mxfp4_dequant + + gate_weight = mxfp4_dequant(tensors[0].cuda(), tensors[1].cuda()).float() + down_weight = mxfp4_dequant(tensors[3].cuda(), tensors[4].cuda()).float() + golden_gate_up = ( + hidden_states.float() @ gate_weight.t() + tensors[2].cuda().float() + ) + golden_gate, golden_up = golden_gate_up[:, ::2], golden_gate_up[:, 1::2] + golden_activated = (golden_up.clamp(-7, 7) + 1) * ( + golden_gate.clamp(max=7) + * torch.sigmoid(golden_gate.clamp(max=7) * 1.702) + ) + golden = golden_activated @ down_weight.t() + tensors[5].cuda().float() + envelope = ( + 8 + * (2**-8) + * ( + golden_activated.abs() @ down_weight.abs().t() + + tensors[5].cuda().float().abs() + ) + + 1e-2 + ) + + assert ((actual.to(golden.device) - golden).abs() <= envelope).all() + assert ((resident.float() - golden).abs() <= envelope).all() + + +@pytest.mark.gpu +def test_native_dispatcher_rejects_mismatched_hidden_dtype(tmp_path): + actual, _, _, _ = _native_dispatch(tmp_path, torch.float32) + + assert torch.count_nonzero(actual).item() == 0 diff --git a/tests/test_gpt_oss_offload_policy.py b/tests/test_gpt_oss_offload_policy.py new file mode 100644 index 00000000..b2fe020a --- /dev/null +++ b/tests/test_gpt_oss_offload_policy.py @@ -0,0 +1,64 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +from moe_infinity.runtime.model_offload import ( + _gpt_oss_offload_enabled, + _wire_expert_collaborators, +) + + +def _engine_config(ratio): + return SimpleNamespace(device_memory_ratio=ratio) + + +def test_gpt_oss_default_ratio_keeps_resident_fallback(): + assert not _gpt_oss_offload_enabled("gpt_oss", _engine_config(0.9)) + + +def test_gpt_oss_low_ratio_enables_dispatcher(): + assert _gpt_oss_offload_enabled("gpt_oss", _engine_config(0.5)) + + +def test_policy_does_not_disable_other_architectures(): + assert _gpt_oss_offload_enabled("mixtral", _engine_config(0.9)) + + +def test_low_ratio_wires_all_gpt_oss_collaborators(): + module = SimpleNamespace() + collaborators = { + "expert_executor": Mock(), + "expert_prefetcher": Mock(), + "expert_tracer": Mock(), + "expert_predictor": Mock(), + "expert_tensor_map": {(0, 0): 7}, + } + + _wire_expert_collaborators(module, True, **collaborators) + + for name, value in collaborators.items(): + assert getattr(module, name) is value + + +def test_resident_mode_leaves_gpt_oss_collaborators_unset(): + module = SimpleNamespace( + expert_executor=None, + expert_prefetcher=None, + expert_tracer=None, + expert_predictor=None, + expert_tensor_map=None, + ) + _wire_expert_collaborators( + module, + False, + expert_executor=Mock(), + expert_prefetcher=Mock(), + expert_tracer=Mock(), + expert_predictor=Mock(), + expert_tensor_map={(0, 0): 7}, + ) + + assert module.expert_executor is None + assert module.expert_prefetcher is None + assert module.expert_tracer is None + assert module.expert_predictor is None + assert module.expert_tensor_map is None diff --git a/tests/test_gpt_oss_wrapper.py b/tests/test_gpt_oss_wrapper.py index ae4f285e..c7a798b0 100644 --- a/tests/test_gpt_oss_wrapper.py +++ b/tests/test_gpt_oss_wrapper.py @@ -1,7 +1,7 @@ import importlib.machinery import sys import types -from unittest.mock import MagicMock +from unittest.mock import MagicMock, Mock import torch @@ -34,6 +34,18 @@ def make_gpt_oss_config(): return cfg +def _deterministic_mlp(): + from moe_infinity.models.gpt_oss import SyncGptOssMLP + + torch.manual_seed(137) + mlp = SyncGptOssMLP(make_gpt_oss_config()) + for parameter in mlp.parameters(): + if parameter.numel(): + torch.nn.init.normal_(parameter, std=0.01) + mlp.layer_id = 1 + return mlp + + def test_sync_gpt_oss_mlp_instantiation(): from moe_infinity.models.gpt_oss import SyncGptOssMLP @@ -104,3 +116,59 @@ def test_sync_gpt_oss_mlp_router_has_bias(): assert mlp.router.bias.shape == ( config.num_local_experts, ), f"Router bias shape must be ({config.num_local_experts},)" + + +def test_sync_gpt_oss_mlp_dispatches_when_executor_is_injected(): + mlp = _deterministic_mlp() + hidden = torch.randn(1, 2, 64) + expected = torch.randn(2, 64) + executor = Mock() + executor.wait_dispatch_local.return_value = expected + mlp.expert_executor = executor + + output, router_logits = mlp(hidden) + + assert output.shape == hidden.shape + assert torch.equal(output.view(2, 64), expected) + executor.dispatch_local.assert_called_once() + args = executor.dispatch_local.call_args + assert args.args[0] == 1 + assert args.args[1].shape == (2, 64) + assert args.args[2].shape == (2, 4) + assert args.args[3].shape == (2, 4) + assert torch.equal(args.kwargs["router_logits"], router_logits) + executor.wait_dispatch_local.assert_called_once_with() + + +def test_sync_gpt_oss_mlp_resident_path_matches_existing_forward(): + mlp = _deterministic_mlp() + hidden = torch.randn(1, 2, 64) + + with torch.no_grad(): + output, router_logits = mlp(hidden) + hidden_flat = hidden.view(-1, hidden.shape[-1]) + routing_weights = torch.softmax( + router_logits, dim=-1, dtype=torch.float32 + ) + routing_weights, selected_experts = torch.topk( + routing_weights, mlp.top_k, dim=-1 + ) + routing_weights = routing_weights / routing_weights.sum( + dim=-1, keepdim=True + ) + expected = torch.zeros_like(hidden_flat) + for expert_idx in range(mlp.num_experts): + token_positions, topk_positions = torch.where( + selected_experts == expert_idx + ) + if token_positions.numel() == 0: + continue + expert_output = mlp._expert_forward( + hidden_flat[token_positions], expert_idx + ) + weights = routing_weights[token_positions, topk_positions].to( + hidden.dtype + ) + expected[token_positions] += weights.unsqueeze(-1) * expert_output + + assert torch.equal(output, expected.view_as(hidden))