Fix NoneType.shape error: device handling, CPU optimizer, safe tensor ops

train_arithmetic.py

@@ -97,21 +97,32 @@ def test_base_model(model, tokenizer, n_samples=20):
     print("\n" + "="*70)
     print("📊 TESTING BASE MODEL PERFORMANCE")
     print("="*70)
-    
+
     test_samples = generate_arithmetic_samples(n_samples)
     correct = 0
-    
+
     model.eval()
     with torch.no_grad():
         for i, sample in enumerate(test_samples):
-            inputs = tokenizer(sample['prompt'], return_tensors='pt').to(model.device)
+            inputs = tokenizer(sample['prompt'], return_tensors='pt')
+
+            # Handle device placement
+            if hasattr(model, 'device') and model.device is not None:
+                inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
             outputs = model.generate(
                 **inputs,
                 max_new_tokens=20,
                 do_sample=False,
                 temperature=1.0
             )
-            response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+
+            # Safely decode response
+            input_ids = inputs.get('input_ids')
+            if input_ids is not None and hasattr(input_ids, 'shape'):
+                response = tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True)
+            else:
+                response = tokenizer.decode(outputs[0], skip_special_tokens=True)
             
             # Extract answer
             numbers = re.findall(r'-?\d+\.?\d*', response)
@@ -157,7 +168,6 @@ def main():
     model = AutoModelForCausalLM.from_pretrained(
         BASE_MODEL,
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto" if torch.cuda.is_available() else None
     )
     
     # Test base model first
@@ -170,6 +180,7 @@ def main():
     print(f"✅ {len(train_dataset)} training samples\n")
     
     # GRPO Config
+    is_cpu = not torch.cuda.is_available()
     training_args = GRPOConfig(
         output_dir="./outputs",
         max_steps=MAX_STEPS,
@@ -179,12 +190,13 @@ def main():
         beta=0.0,  # No KL penalty for this task
         bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
         fp16=False,
-        gradient_checkpointing=True,
-        optim="adamw_8bit",
+        gradient_checkpointing=not is_cpu,  # Disable on CPU
+        optim="adamw_torch" if is_cpu else "adamw_8bit",  # Use standard optimizer on CPU
         logging_steps=1,
         save_steps=MAX_STEPS,  # Save at end
         push_to_hub=False,  # We'll push manually
         report_to="none",
+        use_cpu=is_cpu,  # Explicitly tell trainer to use CPU
     )
     
     print("🚀 Starting GRPO Training...")