refactor(trainer): 优化trainer 结构

tests/test_callbacks.py

@@ -5,10 +5,20 @@ from khaosz.trainer import *
 
 def test_callback_integration(base_test_env, random_dataset):
     """Test that all callbacks are properly integrated"""
+    schedule_config = CosineScheduleConfig(
+        warmup_steps=10,
+        total_steps=20
+    )
+    
     optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+    scheduler = SchedulerFactory.load(optimizer, schedule_config)
+    
     train_config = TrainConfig(
+        model=base_test_env["model"],
+        strategy='seq',
         dataset=random_dataset,
         optimizer=optimizer,
+        scheduler=scheduler,
         checkpoint_dir=base_test_env["test_dir"],
         n_epoch=1,
         batch_size=2,
@@ -18,36 +28,26 @@ def test_callback_integration(base_test_env, random_dataset):
         random_seed=42
     )
     
-    schedule_config = CosineScheduleConfig(
-        warmup_steps=10,
-        total_steps=20
-    )
+
     
     # Create custom callbacks to track calls
     callback_calls = []
     
     class TrackingCallback(TrainCallback):
-        def on_train_begin(self, trainer, context):
+        def on_train_begin(self, context):
             callback_calls.append('on_train_begin')
         
-        def on_batch_end(self, trainer, context):
+        def on_batch_end(self, context):
             callback_calls.append('on_batch_end')
         
-        def on_epoch_end(self, trainer, context):
+        def on_epoch_end(self, context):
             callback_calls.append('on_epoch_end')
     
-    train_config.strategy = StrategyFactory.load(base_test_env["model"], "seq", base_test_env["device"])
-    model_parameter = ModelParameter(
-        base_test_env["model"], 
-        base_test_env["tokenizer"], 
-        base_test_env["transformer_config"]
-    )
+
     
     trainer = Trainer(
-        model_parameter, 
         train_config, 
-        schedule_config,
-        callbacks=[TrackingCallback(), ProgressBarCallback()]
+        callbacks=[TrackingCallback()]
     )
     
     trainer.train()

tests/test_early_stopping.py

@@ -7,35 +7,34 @@ from khaosz.trainer import *
 def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
     """Simulate early stopping behavior"""
     
+    schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
     optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+    scheduler = SchedulerFactory.load(optimizer, schedule_config)
+    
     train_config = TrainConfig(
+        strategy="seq",
+        scheduler=scheduler,
+        model=base_test_env["model"],
         dataset=early_stopping_dataset,
         optimizer=optimizer,
         checkpoint_dir=base_test_env["test_dir"],
         n_epoch=2,
         batch_size=2,
-        checkpoint_interval=2,
+        checkpoint_interval=1,
         accumulation_steps=2,
         random_seed=np.random.randint(1e4),
     )
-    
-    train_config.strategy = StrategyFactory.load(base_test_env["model"], "seq", base_test_env["device"])
-    model_parameter = ModelParameter(
-        base_test_env["model"], 
-        base_test_env["tokenizer"], 
-        base_test_env["transformer_config"]
-    )
-    schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
-    trainer = Trainer(model_parameter, train_config, schedule_config)
+
+    trainer = Trainer(train_config)
     
     # Should handle early stopping gracefully
     checkpoint = None
     try:
         checkpoint = trainer.train()
-        assert len(checkpoint.loss_list) == 2
+        assert checkpoint.iteration == 2
     except Exception:
         # Handle any exceptions
         pass
     
     checkpoint = trainer.train(checkpoint)
-    assert len(checkpoint.loss_list) == 10
+    assert checkpoint.iteration == 10

tests/test_module.py

@@ -51,13 +51,6 @@ def test_env(request: pytest.FixtureRequest):
 
     shutil.rmtree(test_dir)
 
-# parameter loader
-def test_parameter_loader(test_env):
-    loaded_param = ParameterLoader.load(test_env["test_dir"])
-    assert loaded_param.model is not None
-    assert loaded_param.tokenizer is not None
-    assert loaded_param.config == test_env["transformer_config"]
-
 def test_model_parameter(test_env):
     save_dir = os.path.join(test_env["test_dir"], "save")
     model_param = ModelParameter(test_env["model"],test_env["tokenizer"] , test_env["transformer_config"])

tests/test_train_config.py

@@ -31,10 +31,18 @@ def test_gradient_accumulation(base_test_env, random_dataset):
     accumulation_steps_list = [1, 2, 4]
     
     for accumulation_steps in accumulation_steps_list:
+        schedule_config = CosineScheduleConfig(
+            warmup_steps=10,
+            total_steps=20
+        )
         optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+        scheduler = SchedulerFactory.load(optimizer, schedule_config)
         train_config = TrainConfig(
-            dataset=random_dataset,
+            strategy="seq",
+            model=base_test_env["model"],
             optimizer=optimizer,
+            scheduler=scheduler,
+            dataset=random_dataset,
             checkpoint_dir=base_test_env["test_dir"],
             n_epoch=1,
             batch_size=2,
@@ -44,18 +52,7 @@ def test_gradient_accumulation(base_test_env, random_dataset):
             random_seed=42
         )
         
-        schedule_config = CosineScheduleConfig(
-            warmup_steps=10,
-            total_steps=20
-        )
-        train_config.strategy = StrategyFactory.load(base_test_env["model"], "seq", base_test_env["device"])
-        model_parameter = ModelParameter(
-            base_test_env["model"], 
-            base_test_env["tokenizer"], 
-            base_test_env["transformer_config"]
-        )
-        
-        trainer = Trainer(model_parameter, train_config, schedule_config)
+        trainer = Trainer(train_config)
         trainer.train()
         
         assert train_config.accumulation_steps == accumulation_steps

tests/test_train_strategy.py

@@ -35,7 +35,7 @@ def test_schedule_factory_random_configs():
             config.validate()
             
             # Create scheduler using factory
-            scheduler = SchedulerFactory.load_scheduler(optimizer, config)
+            scheduler = SchedulerFactory.load(optimizer, config)
             
             # Verify scheduler type
             if isinstance(config, CosineScheduleConfig):
@@ -83,7 +83,7 @@ def test_schedule_factory_edge_cases():
     
     for config in edge_cases:
         config.validate()
-        scheduler = SchedulerFactory.load_scheduler(optimizer, config)
+        scheduler = SchedulerFactory.load(optimizer, config)
         assert scheduler is not None
         
         # Test multiple steps
@@ -97,16 +97,17 @@ def test_schedule_factory_invalid_configs():
     # Test invalid configurations that should raise errors
     invalid_configs = [
         # Negative warmup steps
-        CosineScheduleConfig(warmup_steps=-10, total_steps=1000, min_rate=0.1),
+        {"warmup_steps": -10, "total_steps": 1000, "min_rate": 0.1},
         # Total steps less than warmup steps
-        CosineScheduleConfig(warmup_steps=500, total_steps=400, min_rate=0.1),
+        {"warmup_steps": 500, "total_steps": 400, "min_rate": 0.1},
         # Invalid min_rate
-        CosineScheduleConfig(warmup_steps=100, total_steps=1000, min_rate=-0.1),
-        CosineScheduleConfig(warmup_steps=100, total_steps=1000, min_rate=1.1),
+        {"warmup_steps": 100, "total_steps": 1000, "min_rate": -0.1},
+        {"warmup_steps": 100, "total_steps": 1000, "min_rate": 1.1},
     ]
     
-    for config in invalid_configs:
+    for kwargs in invalid_configs:
         with pytest.raises(ValueError):
+            config = CosineScheduleConfig(**kwargs)
             config.validate()
 
 
@@ -117,7 +118,7 @@ def test_schedule_factory_state_persistence():
     optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)
     
     config = CosineScheduleConfig(warmup_steps=100, total_steps=1000, min_rate=0.1)
-    scheduler = SchedulerFactory.load_scheduler(optimizer, config)
+    scheduler = SchedulerFactory.load(optimizer, config)
     
     # Take a few steps
     for _ in range(5):
@@ -127,7 +128,7 @@ def test_schedule_factory_state_persistence():
     state_dict = scheduler.state_dict()
     
     # Create new scheduler and load state
-    new_scheduler = SchedulerFactory.load_scheduler(optimizer, config)
+    new_scheduler = SchedulerFactory.load(optimizer, config)
     new_scheduler.load_state_dict(state_dict)
     
     # Verify states match