adds quantization configuration correctly

config/train_gpt_oss_memory_optimized.py

@@ -89,11 +89,8 @@ class GPTOSSMemoryOptimizedConfig:
         
         if self.quantization_config is None:
             self.quantization_config = {
-                "dequantize": True,
-                "load_in_4bit": True,
-                "bnb_4bit_compute_dtype": "bfloat16",
-                "bnb_4bit_use_double_quant": True,
-                "bnb_4bit_quant_type": "nf4"
+                "dequantize": True,  # Use Mxfp4Config as per tutorial
+                "load_in_4bit": False  # Only use 4-bit if explicitly needed
             }
         
         if self.model_kwargs is None:

requirements/requirements_core.txt

@@ -1,6 +1,6 @@
 # Core dependencies for SmolLM3 and GPT-OSS fine-tuning
 torch>=2.0.0
-transformers>=4.55.0  # Updated for GPT-OSS compatibility
+transformers @ git+https://github.com/huggingface/transformers.git  # Latest version with GPT-OSS support
 datasets>=2.14.0
 accelerate>=0.20.0
 peft>=0.17.0  # Updated for GPT-OSS LoRA support

scripts/training/train_gpt_oss.py

@@ -27,26 +27,38 @@ def load_gpt_oss_model_and_tokenizer(config):
     
     # Set up quantization config based on config
     if config.quantization_config and config.quantization_config.get("load_in_4bit"):
-        # Use BitsAndBytesConfig for 4-bit quantization
+        # Use BitsAndBytesConfig for 4-bit quantization (memory optimized)
         quantization_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_compute_dtype=torch.bfloat16,
             bnb_4bit_use_double_quant=True,
             bnb_4bit_quant_type="nf4"
         )
+    elif config.quantization_config and config.quantization_config.get("dequantize"):
+        # Try to use Mxfp4Config if available (as per tutorial)
+        try:
+            from transformers import Mxfp4Config
+            quantization_config = Mxfp4Config(dequantize=True)
+        except ImportError:
+            # Fallback to no quantization if Mxfp4Config not available
+            print("Warning: Mxfp4Config not available, using no quantization")
+            quantization_config = None
     else:
-        # Use BitsAndBytesConfig as default (no quantization)
+        # No quantization
         quantization_config = None
     
     # Model kwargs as per tutorial
     model_kwargs = {
         "attn_implementation": "eager",
         "torch_dtype": torch.bfloat16,
-        "quantization_config": quantization_config,
         "use_cache": False,
         "device_map": "auto",
     }
     
+    # Only add quantization_config if it's not None
+    if quantization_config is not None:
+        model_kwargs["quantization_config"] = quantization_config
+    
     model = AutoModelForCausalLM.from_pretrained(config.model_name, **model_kwargs)
     
     return model, tokenizer