Upload convert_to_linear.py with huggingface_hub

convert_to_linear.py

@@ -109,3 +109,23 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.bfloat16
 )
 
+
+# or to convert on the go, use
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+model = AutoModelForCausalLM.from_pretrained(
+    "openai/gpt-oss-20b",
+    device_map='cuda:0', # modify appropriately.
+    torch_dtype=torch.bfloat16
+)
+
+from tqdm import tqdm
+for layer in tqdm(model.model.layers):
+    experts = layer.mlp.experts
+    if isinstance(experts, GptOssExperts):
+        new_experts = convert_to_linear_experts(experts, model.config) # function is defined in the file above
+        layer.mlp.experts = new_experts.to(model.device, model.dtype)
+print('✅ All experts converted to linear')
+