slightly account for offloading

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/llmcompressor/modeling/gpt_oss.py

@@ -7,6 +7,8 @@
 from transformers.models.gpt_oss.configuration_gpt_oss import GptOssConfig
 from llmcompressor.utils.dev import skip_weights_initialize
 
+from compressed_tensors import update_offload_parameter
+
 
 class GptOssExpert(torch.nn.Module):
     def __init__(self, hidden_size: int, expert_dim: int, alpha: float, limit: float):
@@ -56,18 +58,42 @@ def __init__(self, experts: GptOssExpert):
 
     def load_weights(self, experts: GptOssExperts):
         for expert_index, expert in enumerate(self.experts):
-            expert.gate_proj.weight.data = experts.gate_up_proj[expert_index, ..., ::2].data.T
-            expert.gate_proj.bias.data = experts.gate_up_proj_bias[expert_index, ..., ::2].data
-
-            expert.up_proj.weight.data = experts.gate_up_proj[expert_index, ..., 1::2].data.T
-            expert.up_proj.bias.data = experts.gate_up_proj_bias[expert_index, ..., 1::2].data
+            update_offload_parameter(expert.gate_proj, "weight", experts.gate_up_proj[expert_index, ..., ::2].T)
+            update_offload_parameter(expert.gate_proj, "bias", experts.gate_up_proj_bias[expert_index, ..., ::2])
 
-            expert.down_proj.weight.data = experts.down_proj[expert_index].T
-            expert.down_proj.bias.data = experts.down_proj_bias[expert_index]
+            update_offload_parameter(expert.up_proj, "weight", experts.gate_up_proj[expert_index, ..., 1::2].T)
+            update_offload_parameter(expert.up_proj, "bias", experts.gate_up_proj_bias[expert_index, ..., 1::2])
 
+            update_offload_parameter(expert.down_proj, "weight", experts.down_proj[expert_index].T)
+            update_offload_parameter(expert.down_proj, "bias", experts.down_proj_bias[expert_index])
 
     def to_original(self) -> GptOssExperts:
-        pass
+        with skip_weights_initialize():
+            fake_config = GptOssConfig(
+                intermediate_size=self.intermediate_size,
+                num_local_experts=self.num_experts,
+                hidden_size=self.hidden_size,
+
+            )
+            experts = GptOssExperts(fake_config)
+
+        for expert_index, expert in enumerate(self.experts):
+            experts.gate_up_proj[expert_index, ..., ::2].data = expert.gate_proj.weight.data.T
+            experts.gate_up_proj_bias[expert_index, ..., ::2].data = expert.gate_proj.bias.data
+
+            experts.gate_up_proj[expert_index, ..., 1::2].data = expert.up_proj.weight.data.T
+            experts.gate_up_proj_bias[expert_index, ..., 1::2].data = expert.up_proj.bias.data
+
+            experts.down_proj[expert_index].data = expert.down_proj.weight.data.T
+            experts.down_proj_bias[expert_index] = expert.down_proj.bias.data
+
+        # update offloaded state dict
+        update_offload_parameter(experts, "gate_up_proj", experts.gate_up_proj)
+        update_offload_parameter(experts, "gate_up_proj_bias", experts.gate_up_proj_bias)
+        update_offload_parameter(experts, "down_proj", experts.down_proj)
+        update_offload_parameter(experts, "down_proj_bias", experts.down_proj_bias)
+
+        return experts
 
 
     def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor:
@@ -113,5 +139,8 @@ def forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weig
         linear = GptOssExpertsLinear(original)
         output = linear(input, routing_weights=routing_weights)
 
-        breakpoint()
-        assert torch.allclose(output, true_output, atol=1e-3, rtol=0.0)
+        assert torch.allclose(output, true_output, atol=1e-3, rtol=0.0)
+
+        restored = linear.to_original()
+        restored_output = linear(input, routing_weights=routing_weights)
+        assert torch.allclose(restored_output, true_output, atol=1e-3, rtol=0.0)

src/llmcompressor/modeling/prepare.py

@@ -1,9 +1,11 @@
+import contextlib
 from compressed_tensors.utils import replace_module
 from transformers import PreTrainedModel
 
 from llmcompressor.modeling.deepseek_v3 import replace as replace_deepseekv3
 from llmcompressor.modeling.llama4 import replace as replace_llama4
 from llmcompressor.modeling.qwen3_moe import replace as replace_Qwen3MoE
+from llmcompressor.modeling.gpt_oss import GptOssExpertsLinear
 from llmcompressor.utils.helpers import patch_attr
 
 __all__ = ["replace_modules_for_calibration"]
@@ -42,13 +44,29 @@ def update_qwen3_moe(model, stack):
             )
 
 
-def update_gpt_oss_moe(model, stack):
-    
+def update_gpt_oss_moe(model: PreTrainedModel, stack):
+    @contextlib.contextmanager
+    def replace_context(model, name, module):
+        linear = GptOssExpertsLinear(module)
+        replace_module(model, name, linear)
+        del module
+
+        yield
+
+        restored = linear.to_original()
+        replace_module(model, name, restored)
+
+    # TODO: need to think about duplicates
+    for name, module in model.named_modules():
+        cls_name = module.__class__.__name__
+        if cls_name == "GptOssExpert":
+            stack.enter_context(replace_context(model, name, module))
 
 
 
 moe_context = {
     "Qwen3MoeForCausalLM": update_qwen3_moe,
+    "GptOssForCausalLM": update_gpt_oss_moe,
 }