refactor: 训练循环改为两重迭代并统一参数命名
- 训练循环从三重(epoch→batched→batch)改为二重(epoch→batch) - batch_size → batch_per_device, accumulation_steps → grad_accum_steps - scheduler 移入 step block 对齐 optimizer 更新步 - GradientClippingCallback 改用 on_step_begin 避免零梯度裁剪 - 移除 _train_impl 误导性的 -> Checkpoint 标注 - total_steps 修除为向下取整并精简为一行 - warmup_steps 改为 warmup_ratio (默认0.05)
This commit is contained in:
@@ -31,8 +31,8 @@ def create_train_config(
|
||||
device: str,
|
||||
strategy: str = "seq",
|
||||
n_epoch: int = 1,
|
||||
batch_size: int = 2,
|
||||
accumulation_steps: int = 1,
|
||||
batch_per_device: int = 2,
|
||||
grad_accum_steps: int = 1,
|
||||
max_grad_norm: float = 1.0,
|
||||
ckpt_interval: int = 5,
|
||||
random_seed: int = 42,
|
||||
@@ -47,8 +47,8 @@ def create_train_config(
|
||||
device: Device type ("cuda" or "cpu")
|
||||
strategy: Training strategy type (default: "seq")
|
||||
n_epoch: Number of epochs (default: 1)
|
||||
batch_size: Batch size (default: 2)
|
||||
accumulation_steps: Gradient accumulation steps (default: 1)
|
||||
batch_per_device: Batch size per device (default: 2)
|
||||
grad_accum_steps: Gradient accumulation steps (default: 1)
|
||||
max_grad_norm: Maximum gradient norm for clipping (default: 1.0)
|
||||
ckpt_interval: Checkpoint save interval in iterations (default: 5)
|
||||
random_seed: Random seed for reproducibility (default: 42)
|
||||
@@ -74,9 +74,9 @@ def create_train_config(
|
||||
scheduler_fn=scheduler_fn,
|
||||
ckpt_dir=test_dir,
|
||||
n_epoch=n_epoch,
|
||||
batch_size=batch_size,
|
||||
batch_per_device=batch_per_device,
|
||||
ckpt_interval=ckpt_interval,
|
||||
accumulation_steps=accumulation_steps,
|
||||
grad_accum_steps=grad_accum_steps,
|
||||
max_grad_norm=max_grad_norm,
|
||||
random_seed=random_seed,
|
||||
device_type=device,
|
||||
|
||||
@@ -25,9 +25,9 @@ def test_callback_integration(base_test_env, random_dataset):
|
||||
scheduler_fn=scheduler_fn,
|
||||
ckpt_dir=base_test_env["test_dir"],
|
||||
n_epoch=1,
|
||||
batch_size=2,
|
||||
batch_per_device=2,
|
||||
ckpt_interval=3,
|
||||
accumulation_steps=1,
|
||||
grad_accum_steps=1,
|
||||
max_grad_norm=1.0,
|
||||
random_seed=42,
|
||||
device_type=base_test_env["device"],
|
||||
|
||||
@@ -28,9 +28,9 @@ def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
|
||||
dataset=early_stopping_dataset,
|
||||
ckpt_dir=base_test_env["test_dir"],
|
||||
n_epoch=2,
|
||||
batch_size=2,
|
||||
batch_per_device=2,
|
||||
ckpt_interval=1,
|
||||
accumulation_steps=2,
|
||||
grad_accum_steps=2,
|
||||
random_seed=np.random.randint(1e4),
|
||||
device_type=base_test_env["device"],
|
||||
)
|
||||
|
||||
@@ -7,45 +7,45 @@ def test_different_batch_sizes(base_test_env, random_dataset, train_config_facto
|
||||
"""Test training with different batch sizes"""
|
||||
batch_sizes = [1, 2, 4, 8]
|
||||
|
||||
for batch_size in batch_sizes:
|
||||
for batch_per_device in batch_sizes:
|
||||
train_config = train_config_factory(
|
||||
model=base_test_env["model"],
|
||||
dataset=random_dataset,
|
||||
test_dir=base_test_env["test_dir"],
|
||||
device=base_test_env["device"],
|
||||
batch_size=batch_size,
|
||||
batch_per_device=batch_per_device,
|
||||
)
|
||||
|
||||
assert train_config.batch_size == batch_size
|
||||
assert train_config.batch_per_device == batch_per_device
|
||||
|
||||
|
||||
def test_gradient_accumulation(base_test_env, random_dataset, train_config_factory):
|
||||
"""Test training with different gradient accumulation steps"""
|
||||
accumulation_steps_list = [1, 2, 4]
|
||||
grad_accum_steps_list = [1, 2, 4]
|
||||
|
||||
for accumulation_steps in accumulation_steps_list:
|
||||
for grad_accum_steps in grad_accum_steps_list:
|
||||
train_config = train_config_factory(
|
||||
model=base_test_env["model"],
|
||||
dataset=random_dataset,
|
||||
test_dir=base_test_env["test_dir"],
|
||||
device=base_test_env["device"],
|
||||
batch_size=2,
|
||||
accumulation_steps=accumulation_steps,
|
||||
batch_per_device=2,
|
||||
grad_accum_steps=grad_accum_steps,
|
||||
)
|
||||
|
||||
trainer = Trainer(train_config)
|
||||
trainer.train()
|
||||
|
||||
assert train_config.accumulation_steps == accumulation_steps
|
||||
assert train_config.grad_accum_steps == grad_accum_steps
|
||||
|
||||
|
||||
def test_memory_efficient_training(base_test_env, random_dataset, train_config_factory):
|
||||
"""Test training with memory-efficient configurations"""
|
||||
# Test with smaller batch sizes and gradient checkpointing
|
||||
small_batch_configs = [
|
||||
{"batch_size": 1, "accumulation_steps": 8},
|
||||
{"batch_size": 2, "accumulation_steps": 4},
|
||||
{"batch_size": 4, "accumulation_steps": 2},
|
||||
{"batch_per_device": 1, "grad_accum_steps": 8},
|
||||
{"batch_per_device": 2, "grad_accum_steps": 4},
|
||||
{"batch_per_device": 4, "grad_accum_steps": 2},
|
||||
]
|
||||
|
||||
for config in small_batch_configs:
|
||||
@@ -54,8 +54,9 @@ def test_memory_efficient_training(base_test_env, random_dataset, train_config_f
|
||||
dataset=random_dataset,
|
||||
test_dir=base_test_env["test_dir"],
|
||||
device=base_test_env["device"],
|
||||
batch_size=config["batch_size"],
|
||||
accumulation_steps=config["accumulation_steps"],
|
||||
batch_per_device=config["batch_per_device"],
|
||||
grad_accum_steps=config["grad_accum_steps"],
|
||||
)
|
||||
|
||||
assert train_config.accumulation_steps == config["accumulation_steps"]
|
||||
assert train_config.grad_accum_steps == config["grad_accum_steps"]
|
||||
assert train_config.batch_per_device == config["batch_per_device"]
|
||||
|
||||
Reference in New Issue
Block a user