From 435b4f16e6326416161d4f39408d8971643c53d6 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Thu, 6 Aug 2026 13:58:29 +0000 Subject: [PATCH 1/6] fix(gpt-oss): resident-load MXFP4 experts + sinks + router GPT-OSS is excluded from the C++ expert dispatcher and runs a resident Python expert loop in SyncGptOssMLP.forward, but the loader never materialized the _PackedExperts params: expert weights stayed zeros, biases NaN, attention sinks and router garbage, producing incoherent output. Add _load_resident_gpt_oss to load the real MXFP4 blocks/scales (uint8, output-major [E,N,K//2]/[E,N,K//32]), biases, router, and self_attn.sinks into the live params and drop them from name_id_map. Fix _expert_forward_mxfp4 to feed packed weights without the erroneous .t() (checkpoint is already output-major; verified rel-err 0.0 vs reference dequant on gate_up and down proj). Real 120B: base output now coherent ('The capital of France is Paris.'); DFlash agreement 0.95->1.00, mean acceptance 1.0->6.57, ~21x decode speedup. 103 dflash tests pass. --- moe_infinity/models/gpt_oss.py | 13 +-- moe_infinity/runtime/model_offload.py | 143 ++++++++++++++++++++++++++ 2 files changed, 150 insertions(+), 6 deletions(-) diff --git a/moe_infinity/models/gpt_oss.py b/moe_infinity/models/gpt_oss.py index b36ab65b..24d6c3a9 100644 --- a/moe_infinity/models/gpt_oss.py +++ b/moe_infinity/models/gpt_oss.py @@ -99,12 +99,13 @@ def _expert_forward_mxfp4( down_scales = self.experts.down_proj_scales[expert_idx].to(device) down_b = self.experts.down_proj_bias[expert_idx].to(device) - # Checkpoint stores weights as [K, N//2] (input-major packed). - # The fused kernel expects [N, K//2] (output-major packed). + # Checkpoint stores each expert's packed weights output-major as + # [N, K//2] (blocks) and [N, K//32] (scales) - exactly the layout the + # fused kernel expects, so they are fed through unchanged. gate_up_out = fused_mxfp4_gemm( x, - gate_up_packed.t().contiguous(), - gate_up_scales.t().contiguous(), + gate_up_packed.contiguous(), + gate_up_scales.contiguous(), gate_up_b, ) @@ -113,8 +114,8 @@ def _expert_forward_mxfp4( result = fused_mxfp4_gemm( activated.to(torch.bfloat16), - down_packed.t().contiguous(), - down_scales.t().contiguous(), + down_packed.contiguous(), + down_scales.contiguous(), down_b, ) return result diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index 45d2d6b5..c4a63bf9 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -849,6 +849,9 @@ def archer_from_pretrained(cls, *args, **kwargs): module_idx += 1 + if getattr(self.config, "model_type", "") == "gpt_oss": + self._load_resident_gpt_oss(model) + self.setup_archer_hooks(model) return model @@ -961,6 +964,146 @@ def _load_resident_shared_experts(self, model): f"Missing shared_experts weights: {sorted(remaining)[:5]}" ) + @torch.no_grad() + def _load_resident_gpt_oss(self, model): + # GPT-OSS is not wired into the C++ expert dispatcher (see the + # SyncGptOssMLP guards in setup_archer_hooks / register_expert): its + # SyncGptOssMLP.forward runs a resident Python expert loop reading the + # _PackedExperts params directly. The Archer empty-init replaces those + # params with [1] placeholders and never materializes them, so load the + # real MXFP4 blocks/scales, biases, router, and attention sinks here and + # keep them resident. MXFP4 packed weights stay uint8 output-major + # ([E, N, K//2] blocks, [E, N, K//32] scales) - the layout the fused + # kernel consumes without transposition. + modules = [ + m for m in model.modules() if isinstance(m, SyncGptOssMLP) + ] + if not modules: + return + + num_devices = torch.cuda.device_count() + target = get_device(num_devices - 1) if num_devices else "cpu" + + def _set(param, tensor, *, cast_dtype=None): + data = tensor + if cast_dtype is not None: + data = data.to(cast_dtype) + param.requires_grad_(False) + param.data = data.to(device=target).contiguous() + param._moe_infinity_resident = True + + loaded_names = set() + + def _reshape_blocks(t): + if t.dim() == 4: + E, N, G, B = t.shape + return t.reshape(E, N, G * B) + return t + + for module in modules: + layer_id = module.layer_id + prefix = f"model.layers.{layer_id}" + experts = f"{prefix}.mlp.experts" + wanted = { + f"{experts}.gate_up_proj_blocks": ( + module.experts.gate_up_proj, + _reshape_blocks, + None, + ), + f"{experts}.gate_up_proj_scales": ( + module.experts.gate_up_proj_scales, + None, + None, + ), + f"{experts}.gate_up_proj_bias": ( + module.experts.gate_up_proj_bias, + None, + self.dtype_cls, + ), + f"{experts}.down_proj_blocks": ( + module.experts.down_proj, + _reshape_blocks, + None, + ), + f"{experts}.down_proj_scales": ( + module.experts.down_proj_scales, + None, + None, + ), + f"{experts}.down_proj_bias": ( + module.experts.down_proj_bias, + None, + self.dtype_cls, + ), + f"{prefix}.mlp.router.weight": ( + module.router.weight, + None, + self.dtype_cls, + ), + f"{prefix}.mlp.router.bias": ( + module.router.bias, + None, + self.dtype_cls, + ), + } + remaining = dict(wanted) + for ckpt in self.ckpt_files: + if not remaining or not ckpt.endswith(".safetensors"): + continue + with safe_open(ckpt, framework="pt", device="cpu") as f: + keys = set(f.keys()) + for name in list(remaining): + if name not in keys: + continue + param, transform, cast_dtype = remaining[name] + t = f.get_tensor(name) + if transform is not None: + t = transform(t) + _set(param, t, cast_dtype=cast_dtype) + loaded_names.add(name) + del remaining[name] + + if remaining: + raise RuntimeError( + f"Missing gpt-oss resident weights: {sorted(remaining)[:5]}" + ) + + sink_params = { + name: param + for name, param in model.named_parameters(recurse=True) + if name.endswith("self_attn.sinks") + } + sink_remaining = dict(sink_params) + for ckpt in self.ckpt_files: + if not sink_remaining or not ckpt.endswith(".safetensors"): + continue + with safe_open(ckpt, framework="pt", device="cpu") as f: + keys = set(f.keys()) + for name in list(sink_remaining): + lookup = name + if lookup not in keys and name.startswith( + model.base_model_prefix + "." + ): + lookup = name[len(model.base_model_prefix) + 1 :] + if lookup not in keys: + continue + _set( + sink_remaining[name], + f.get_tensor(lookup), + cast_dtype=self.dtype_cls, + ) + loaded_names.add(name) + loaded_names.add(lookup) + del sink_remaining[name] + if sink_remaining: + raise RuntimeError( + f"Missing gpt-oss attention sinks: {sorted(sink_remaining)[:5]}" + ) + + for name in loaded_names: + for candidate in (name, f"{model.base_model_prefix}.{name}"): + self.name_id_map.pop(candidate, None) + def get_topology(self, model): name_lst = [] ret_dict = {} From 54d4013f0c271e944511c004695dc8fcda3aad1c Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Thu, 6 Aug 2026 15:27:35 +0000 Subject: [PATCH 2/6] style(gpt-oss): apply ruff format (v0.6.9) to loader fix Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- moe_infinity/runtime/model_offload.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index c4a63bf9..739f3f48 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -975,9 +975,7 @@ def _load_resident_gpt_oss(self, model): # keep them resident. MXFP4 packed weights stay uint8 output-major # ([E, N, K//2] blocks, [E, N, K//32] scales) - the layout the fused # kernel consumes without transposition. - modules = [ - m for m in model.modules() if isinstance(m, SyncGptOssMLP) - ] + modules = [m for m in model.modules() if isinstance(m, SyncGptOssMLP)] if not modules: return From 64e2d6be1805eb01b43d076541426750c48118b4 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Fri, 14 Aug 2026 09:16:15 +0000 Subject: [PATCH 3/6] feat(gpt-oss): parse per-expert packed tensor identities Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- moe_infinity/utils/hf_config.py | 11 ++++++---- tests/test_gpt_oss_config.py | 38 ++++++++++++++++++--------------- 2 files changed, 28 insertions(+), 21 deletions(-) diff --git a/moe_infinity/utils/hf_config.py b/moe_infinity/utils/hf_config.py index 901aba4a..293c6818 100644 --- a/moe_infinity/utils/hf_config.py +++ b/moe_infinity/utils/hf_config.py @@ -126,7 +126,7 @@ def parse_expert_id( param_name: str, config: PretrainedConfig ) -> Tuple[Optional[int], Optional[int]]: arch = (config.architectures or [""])[0].lower() - _, _, num_encoder_layers = parse_moe_param(config) + num_layers, _, num_encoder_layers = parse_moe_param(config) result = None layer_type = "" layer_id = 0 @@ -196,13 +196,16 @@ def parse_expert_id( layer_id = int(layer_id) expert_id = int(expert_id) elif "gpt_oss" in arch or "gptoss" in arch: + layer_type = "decoder" result = re.findall( - r"layers\.(\d+)\.mlp\.experts\.(gate_up_proj|down_proj)", + r"layers\.(\d+)\.mlp\.experts\.(\d+)\." + r"(?:gate_up_proj|down_proj)_(?:blocks|scales|bias)$", param_name, ) if result: - layer_id = int(result[0][0]) - return layer_id, None + layer_id, expert_id = (int(value) for value in result[0]) + if layer_id >= num_layers or expert_id >= config.num_local_experts: + return None, None if result: if layer_type == "decoder": diff --git a/tests/test_gpt_oss_config.py b/tests/test_gpt_oss_config.py index 5c90fef0..9ed5710e 100644 --- a/tests/test_gpt_oss_config.py +++ b/tests/test_gpt_oss_config.py @@ -37,37 +37,41 @@ def test_parse_moe_param_gpt_oss(): assert num_encoder_layers == 0 -def test_parse_expert_id_gpt_oss_packed(): +def test_parse_expert_id_gpt_oss_gate_up_slice(): from moe_infinity.utils.hf_config import parse_expert_id config = make_gpt_oss_config() - layer_id, expert_id = parse_expert_id( - "model.layers.5.mlp.experts.gate_up_proj_blocks", config - ) - assert layer_id == 5 - assert expert_id is None + assert parse_expert_id( + "model.layers.5.mlp.experts.17.gate_up_proj_blocks", config + ) == (5, 17) -def test_parse_expert_id_gpt_oss_router(): +def test_parse_expert_id_gpt_oss_down_slice(): from moe_infinity.utils.hf_config import parse_expert_id config = make_gpt_oss_config() - layer_id, expert_id = parse_expert_id( - "model.layers.5.mlp.router.weight", config - ) - assert layer_id is None - assert expert_id is None + assert parse_expert_id( + "model.layers.11.mlp.experts.31.down_proj_bias", config + ) == (11, 31) -def test_parse_expert_id_gpt_oss_down_proj(): +def test_parse_expert_id_gpt_oss_rejects_out_of_range_expert(): + from moe_infinity.utils.hf_config import parse_expert_id + + config = make_gpt_oss_config() + assert parse_expert_id( + "model.layers.5.mlp.experts.32.gate_up_proj_blocks", config + ) == (None, None) + + +def test_parse_expert_id_gpt_oss_router(): from moe_infinity.utils.hf_config import parse_expert_id config = make_gpt_oss_config() - layer_id, expert_id = parse_expert_id( - "model.layers.11.mlp.experts.down_proj_blocks", config + assert parse_expert_id("model.layers.5.mlp.router.weight", config) == ( + None, + None, ) - assert layer_id == 11 - assert expert_id is None def test_parse_expert_dtype_gpt_oss_none(): From 93f01335529536a389e6ac223b2894279f4a0d0e Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Fri, 14 Aug 2026 09:17:17 +0000 Subject: [PATCH 4/6] feat(gpt-oss): split packed experts into offload views Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- moe_infinity/runtime/model_offload.py | 58 +++++++++++++++++ tests/test_gpt_oss_offload_topology.py | 86 ++++++++++++++++++++++++++ 2 files changed, 144 insertions(+) create mode 100644 tests/test_gpt_oss_offload_topology.py diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index 739f3f48..adaf5edc 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -164,6 +164,60 @@ def _out_block(experts_prefix: str) -> str: del state_dict[down_key] +_GPT_OSS_EXPERT_FIELDS = ( + "gate_up_proj_blocks", + "gate_up_proj_scales", + "gate_up_proj_bias", + "down_proj_blocks", + "down_proj_scales", + "down_proj_bias", +) + + +def _expand_gpt_oss_packed_experts(state_dict, config): + if getattr(config, "model_type", "") != "gpt_oss": + return + + layer_prefixes = sorted( + { + key.rsplit(".", 1)[0] + for key in state_dict + if key.endswith(".mlp.experts.gate_up_proj_blocks") + } + ) + expected_experts = int(config.num_local_experts) + for prefix in layer_prefixes: + packed = {} + missing = [] + for field in _GPT_OSS_EXPERT_FIELDS: + key = f"{prefix}.{field}" + if key not in state_dict: + missing.append(field) + else: + packed[field] = state_dict[key] + if missing: + raise ValueError( + f"Incomplete GPT-OSS packed expert layer {prefix}: {missing}" + ) + for field, tensor in packed.items(): + if tensor.shape[0] != expected_experts: + raise ValueError( + f"{prefix}.{field} has {tensor.shape[0]} experts; " + f"expected {expected_experts}" + ) + for expert_idx in range(expected_experts): + expert_prefix = f"{prefix}.{expert_idx}" + for field in _GPT_OSS_EXPERT_FIELDS: + view = packed[field][expert_idx] + if not view.is_contiguous(): + raise ValueError( + f"Non-contiguous GPT-OSS slice {expert_prefix}.{field}" + ) + state_dict[f"{expert_prefix}.{field}"] = view + for field in _GPT_OSS_EXPERT_FIELDS: + del state_dict[f"{prefix}.{field}"] + + def _compute_config_fingerprint(config: object) -> str: fields = [ "model_type", @@ -637,6 +691,9 @@ def archer_from_pretrained(cls, *args, **kwargs): state_dict = torch.load(ckpt) _remap_v5_batched_experts(state_dict, self.config) + _expand_gpt_oss_packed_experts( + state_dict, self.config + ) is_gptq_ckpt = is_gptq_quantized(self.config) self._cast_state_dict_tensors( @@ -647,6 +704,7 @@ def archer_from_pretrained(cls, *args, **kwargs): if ( is_mxfp4_ckpt + and self.config.model_type != "gpt_oss" and os.environ.get("MOE_INFINITY_MXFP4_DEQUANT", "") == "1" ): diff --git a/tests/test_gpt_oss_offload_topology.py b/tests/test_gpt_oss_offload_topology.py new file mode 100644 index 00000000..37de7c77 --- /dev/null +++ b/tests/test_gpt_oss_offload_topology.py @@ -0,0 +1,86 @@ +from types import SimpleNamespace + +import torch + +from moe_infinity.runtime.model_offload import ( + _expand_gpt_oss_packed_experts, +) + + +def _config(layers=2, experts=128): + return SimpleNamespace( + architectures=["GptOssForCausalLM"], + model_type="gpt_oss", + num_hidden_layers=layers, + num_local_experts=experts, + ) + + +def _packed_layer(layer_id, experts=128): + prefix = f"model.layers.{layer_id}.mlp.experts" + return { + f"{prefix}.gate_up_proj_blocks": torch.empty( + experts, 12, 8, dtype=torch.uint8 + ), + f"{prefix}.gate_up_proj_scales": torch.empty( + experts, 12, 1, dtype=torch.uint8 + ), + f"{prefix}.gate_up_proj_bias": torch.empty( + experts, 12, dtype=torch.bfloat16 + ), + f"{prefix}.down_proj_blocks": torch.empty( + experts, 6, 6, dtype=torch.uint8 + ), + f"{prefix}.down_proj_scales": torch.empty( + experts, 6, 1, dtype=torch.uint8 + ), + f"{prefix}.down_proj_bias": torch.empty( + experts, 6, dtype=torch.bfloat16 + ), + } + + +def test_expansion_creates_128_identities_per_layer_without_copy(): + state = {**_packed_layer(0), **_packed_layer(1)} + originals = dict(state) + + _expand_gpt_oss_packed_experts(state, _config()) + + prefixes = { + key.rsplit(".", 1)[0] for key in state if ".mlp.experts." in key + } + assert len(prefixes) == 128 * 2 + assert len(state) == 128 * 2 * 6 + + for layer_id in range(2): + packed_prefix = f"model.layers.{layer_id}.mlp.experts" + for expert_idx in range(128): + expert_prefix = f"{packed_prefix}.{expert_idx}" + for field in ( + "gate_up_proj_blocks", + "gate_up_proj_scales", + "gate_up_proj_bias", + "down_proj_blocks", + "down_proj_scales", + "down_proj_bias", + ): + view = state[f"{expert_prefix}.{field}"] + packed = originals[f"{packed_prefix}.{field}"] + assert view.is_contiguous() + assert ( + view.untyped_storage().data_ptr() + == packed.untyped_storage().data_ptr() + ) + assert view.storage_offset() == expert_idx * packed.stride(0) + + +def test_expansion_rejects_incomplete_layer(): + state = _packed_layer(0) + del state["model.layers.0.mlp.experts.down_proj_scales"] + + try: + _expand_gpt_oss_packed_experts(state, _config(layers=1)) + except ValueError as exc: + assert "down_proj_scales" in str(exc) + else: + raise AssertionError("incomplete GPT-OSS packed layer was accepted") From 194e720a3fb86334cae0d05337d69023d3310e39 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Fri, 14 Aug 2026 09:26:03 +0000 Subject: [PATCH 5/6] feat(gpt-oss): register packed slices as Archer experts Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- moe_infinity/runtime/model_offload.py | 64 +++++++++++++++++++++++--- tests/test_gpt_oss_offload_topology.py | 45 ++++++++++++++++++ 2 files changed, 103 insertions(+), 6 deletions(-) diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index adaf5edc..a5f88c5d 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -218,6 +218,53 @@ def _expand_gpt_oss_packed_experts(state_dict, config): del state_dict[f"{prefix}.{field}"] +def _gpt_oss_expert_groups(name_id_map, config): + fields = {name: index for index, name in enumerate(_GPT_OSS_EXPERT_FIELDS)} + grouped = {} + for name, tensor_id in name_id_map.items(): + layer_id, expert_id = parse_expert_id(name, config) + if layer_id is None or expert_id is None: + continue + field = name.rsplit(".", 1)[-1] + if field not in fields: + continue + slots = grouped.setdefault( + (layer_id, expert_id), [None] * len(_GPT_OSS_EXPERT_FIELDS) + ) + slots[fields[field]] = tensor_id + + topology = [] + for layer_id in range(config.num_hidden_layers): + experts = [] + for expert_id in range(config.num_local_experts): + ids = grouped.get((layer_id, expert_id)) + if ids is None or any(tensor_id is None for tensor_id in ids): + raise ValueError( + f"Missing GPT-OSS expert tensors for layer={layer_id}, " + f"expert={expert_id}" + ) + experts.append(ids) + topology.append((f"model.layers.{layer_id}.mlp.experts", experts)) + return topology + + +def _make_expert_tensor_map(name_id_map, config): + if getattr(config, "model_type", "") == "gpt_oss": + return { + (layer_id, expert_id): tensor_ids[0] + for layer_id, (_, experts) in enumerate( + _gpt_oss_expert_groups(name_id_map, config) + ) + for expert_id, tensor_ids in enumerate(experts) + } + result = {} + for name, tensor_id in name_id_map.items(): + layer_id, expert_id = parse_expert_id(name, config) + if expert_id is not None: + result[(layer_id, expert_id)] = tensor_id + return result + + def _compute_config_fingerprint(config: object) -> str: fields = [ "model_type", @@ -841,11 +888,9 @@ def archer_from_pretrained(cls, *args, **kwargs): model.model.encoder.embed_tokens.weight.ar_id = 0 model.model.decoder.embed_tokens.weight.ar_id = 0 - self.expert_tensor_map = dict() - for name, id in self.name_id_map.items(): - layer_id, expert_id = parse_expert_id(name, self.config) - if expert_id is not None: - self.expert_tensor_map[(layer_id, expert_id)] = id + self.expert_tensor_map = _make_expert_tensor_map( + self.name_id_map, self.config + ) self.expert_prefetcher.expert_tensor_map = ( self.expert_tensor_map ) @@ -1164,6 +1209,13 @@ def get_topology(self, model): name_lst = [] ret_dict = {} + if getattr(self.config, "model_type", "") == "gpt_oss": + gpt_oss_topology = _gpt_oss_expert_groups( + self.name_id_map, self.config + ) + name_lst.extend(name for name, _ in gpt_oss_topology) + ret_dict.update(dict(gpt_oss_topology)) + for name, _ in model.named_parameters(recurse=True): match = re.search(r"\d+", name) if name not in self.name_id_map: @@ -1360,7 +1412,7 @@ def gen_args_hook( key = key.split(".")[0] output_device_index = 0 - if "expert" in key and self.config.model_type != "gpt_oss": + if "expert" in key: for expert_idx, expert_tensors in enumerate(tensors): expert_key = ( f"{key}.expert_{expert_idx}" diff --git a/tests/test_gpt_oss_offload_topology.py b/tests/test_gpt_oss_offload_topology.py index 37de7c77..fdaa883b 100644 --- a/tests/test_gpt_oss_offload_topology.py +++ b/tests/test_gpt_oss_offload_topology.py @@ -4,6 +4,8 @@ from moe_infinity.runtime.model_offload import ( _expand_gpt_oss_packed_experts, + _gpt_oss_expert_groups, + _make_expert_tensor_map, ) @@ -84,3 +86,46 @@ def test_expansion_rejects_incomplete_layer(): assert "down_proj_scales" in str(exc) else: raise AssertionError("incomplete GPT-OSS packed layer was accepted") + + +def _synthetic_name_id_map(layers=2, experts=128): + mapping = {} + tensor_id = 100 + for layer_id in range(layers): + for expert_idx in range(experts): + prefix = f"model.layers.{layer_id}.mlp.experts.{expert_idx}" + for field in ( + "gate_up_proj_blocks", + "gate_up_proj_scales", + "gate_up_proj_bias", + "down_proj_blocks", + "down_proj_scales", + "down_proj_bias", + ): + mapping[f"{prefix}.{field}"] = tensor_id + tensor_id += 1 + return mapping + + +def test_gpt_oss_topology_has_six_ordered_ids_for_every_expert(): + config = _config() + name_id_map = _synthetic_name_id_map() + + groups = _gpt_oss_expert_groups(name_id_map, config) + + assert len(groups) == 2 + assert all(len(experts) == 128 for _, experts in groups) + assert all(len(ids) == 6 for _, experts in groups for ids in experts) + first_ids = groups[0][1][0] + assert first_ids == [100, 101, 102, 103, 104, 105] + + +def test_expert_tensor_map_has_128_times_num_layers_entries(): + config = _config() + name_id_map = _synthetic_name_id_map() + + tensor_map = _make_expert_tensor_map(name_id_map, config) + + assert len(tensor_map) == 128 * config.num_hidden_layers + assert tensor_map[(0, 0)] == 100 + assert tensor_map[(1, 127)] == max(name_id_map.values()) - 5 From 67fe9797bd2facbbee2225417be735bc44bf0d1b Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Fri, 14 Aug 2026 19:43:22 +0000 Subject: [PATCH 6/6] style: apply ruff-format and clang-format --- moe_infinity/runtime/model_offload.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index a5f88c5d..3db35d70 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -738,9 +738,7 @@ def archer_from_pretrained(cls, *args, **kwargs): state_dict = torch.load(ckpt) _remap_v5_batched_experts(state_dict, self.config) - _expand_gpt_oss_packed_experts( - state_dict, self.config - ) + _expand_gpt_oss_packed_experts(state_dict, self.config) is_gptq_ckpt = is_gptq_quantized(self.config) self._cast_state_dict_tensors(