attempts to revert to previous training arguments

model.py

@@ -128,50 +128,45 @@ class SmolLM3Model:
         logger.info(f"Config type: {type(self.config)}")
         logger.info(f"Config attributes: {[attr for attr in dir(self.config) if not attr.startswith('_')]}")
         
-        # Merge config with kwargs
-        training_args = {}
-        
-        # Add arguments one by one with error checking
-        try:
-            training_args["output_dir"] = output_dir
-            training_args["per_device_train_batch_size"] = self.config.batch_size
-            training_args["per_device_eval_batch_size"] = self.config.batch_size
-            training_args["gradient_accumulation_steps"] = self.config.gradient_accumulation_steps
-            training_args["learning_rate"] = self.config.learning_rate
-            training_args["weight_decay"] = self.config.weight_decay
-            training_args["warmup_steps"] = self.config.warmup_steps
-            training_args["max_steps"] = self.config.max_iters
-            training_args["save_steps"] = self.config.save_steps
-            training_args["eval_steps"] = self.config.eval_steps
-            training_args["logging_steps"] = self.config.logging_steps
-            training_args["save_total_limit"] = self.config.save_total_limit
-            training_args["eval_strategy"] = self.config.eval_strategy
-            training_args["metric_for_best_model"] = self.config.metric_for_best_model
-            training_args["greater_is_better"] = self.config.greater_is_better
-            training_args["load_best_model_at_end"] = self.config.load_best_model_at_end
-            training_args["fp16"] = self.config.fp16
-            training_args["bf16"] = self.config.bf16
-            training_args["ddp_backend"] = self.config.ddp_backend if torch.cuda.device_count() > 1 else None
-            training_args["ddp_find_unused_parameters"] = self.config.ddp_find_unused_parameters if torch.cuda.device_count() > 1 else False
-            training_args["report_to"] = None
-            training_args["remove_unused_columns"] = False
-            training_args["dataloader_pin_memory"] = False
-            training_args["group_by_length"] = True
-            training_args["length_column_name"] = "length"
-            training_args["ignore_data_skip"] = False
-            training_args["seed"] = 42
-            training_args["data_seed"] = 42
-            training_args["dataloader_num_workers"] = getattr(self.config, 'dataloader_num_workers', 4)
-            training_args["max_grad_norm"] = getattr(self.config, 'max_grad_norm', 1.0)
-            training_args["optim"] = self.config.optimizer
-            training_args["lr_scheduler_type"] = self.config.scheduler
-            training_args["warmup_ratio"] = 0.1
-            training_args["save_strategy"] = "steps"
-            training_args["logging_strategy"] = "steps"
-            training_args["prediction_loss_only"] = True
-        except Exception as e:
-            logger.error(f"Error creating training arguments: {e}")
-            raise
+        # Merge config with kwargs - using the working approach from the functioning commit
+        training_args = {
+            "output_dir": output_dir,
+            "per_device_train_batch_size": self.config.batch_size,
+            "per_device_eval_batch_size": self.config.batch_size,
+            "gradient_accumulation_steps": self.config.gradient_accumulation_steps,
+            "learning_rate": self.config.learning_rate,
+            "weight_decay": self.config.weight_decay,
+            "warmup_steps": self.config.warmup_steps,
+            "max_steps": self.config.max_iters,
+            "save_steps": self.config.save_steps,
+            "eval_steps": self.config.eval_steps,
+            "logging_steps": self.config.logging_steps,
+            "save_total_limit": self.config.save_total_limit,
+            "eval_strategy": self.config.eval_strategy,
+            "metric_for_best_model": self.config.metric_for_best_model,
+            "greater_is_better": self.config.greater_is_better,
+            "load_best_model_at_end": self.config.load_best_model_at_end,
+            "fp16": self.config.fp16,
+            "bf16": self.config.bf16,
+            "ddp_backend": self.config.ddp_backend if torch.cuda.device_count() > 1 else None,
+            "ddp_find_unused_parameters": self.config.ddp_find_unused_parameters if torch.cuda.device_count() > 1 else False,
+            "report_to": None,
+            "remove_unused_columns": False,
+            "dataloader_pin_memory": False,
+            "group_by_length": True,
+            "length_column_name": "length",
+            "ignore_data_skip": False,
+            "seed": 42,
+            "data_seed": 42,
+            "dataloader_num_workers": getattr(self.config, 'dataloader_num_workers', 4),
+            "max_grad_norm": getattr(self.config, 'max_grad_norm', 1.0),
+            "optim": self.config.optimizer,
+            "lr_scheduler_type": self.config.scheduler,
+            "warmup_ratio": 0.1,
+            "save_strategy": "steps",
+            "logging_strategy": "steps",
+            "prediction_loss_only": True,
+        }
         
         # Override with kwargs
         training_args.update(kwargs)

trainer.py

@@ -101,26 +101,25 @@ class SmolLM3Trainer:
                     eval_loss = metrics.get('eval_loss', 'N/A')
                     print(f"📊 Evaluation at step {step}: eval_loss={eval_loss}")
         
-        # Temporarily disable callbacks to debug the issue
+        # Add monitoring callbacks
         callbacks = []
-        logger.info("Callbacks disabled for debugging")
         
-        # # Add simple console callback
-        # callbacks.append(SimpleConsoleCallback())
-        # logger.info("Added simple console monitoring callback")
-        # 
-        # # Try to add Trackio callback if available
-        # if self.monitor and self.monitor.enable_tracking:
-        #     try:
-        #         trackio_callback = self.monitor.create_monitoring_callback()
-        #         if trackio_callback:
-        #         callbacks.append(trackio_callback)
-        #         logger.info("Added Trackio monitoring callback")
-        #     else:
-        #         logger.warning("Failed to create Trackio callback")
-        # except Exception as e:
-        #     logger.error(f"Error creating Trackio callback: {e}")
-        #     logger.info("Continuing with console monitoring only")
+        # Add simple console callback
+        callbacks.append(SimpleConsoleCallback())
+        logger.info("Added simple console monitoring callback")
+        
+        # Try to add Trackio callback if available
+        if self.monitor and self.monitor.enable_tracking:
+            try:
+                trackio_callback = self.monitor.create_monitoring_callback()
+                if trackio_callback:
+                    callbacks.append(trackio_callback)
+                    logger.info("Added Trackio monitoring callback")
+                else:
+                    logger.warning("Failed to create Trackio callback")
+            except Exception as e:
+                logger.error(f"Error creating Trackio callback: {e}")
+                logger.info("Continuing with console monitoring only")
         
         # Try standard Trainer first (more stable with callbacks)
         logger.info("Creating Trainer with training arguments...")