ViperEkura
4145d35e3c
- model: nn.Module -> model_fn 工厂函数，spawn 边界只传字符串 - Trainer.train(resume_dir=path) — Checkpoint 不再通过 pickle 传递 - TrainContextBuilder.with_resume_dir(path) — 自动检测 meta.json 分流 resume/from-scratch - CheckpointCallback: 拆分 state_dict 收集（全 rank）与磁盘写入（rank-0），修复 FSDP 死锁 - serialization: load_torch 支持 broadcast，消除 _load_extra/_load_torch_broadcast - optimizer/scheduler 恢复逻辑内联到 build()，在 executor.prepare() 之后执行 - pyproject.toml: ruff exclude build/ 避免 CI 扫描构建产物
63 lines
2.2 KiB
Python
63 lines
2.2 KiB
Python
from astrai.trainer import Trainer
|
|
|
|
# train_config_factory is injected via fixture
|
|
|
|
|
|
def test_different_batch_sizes(base_test_env, random_dataset, train_config_factory):
|
|
"""Test training with different batch sizes"""
|
|
batch_sizes = [1, 2, 4, 8]
|
|
|
|
for batch_per_device in batch_sizes:
|
|
train_config = train_config_factory(
|
|
model_fn=lambda: base_test_env["model"],
|
|
dataset=random_dataset,
|
|
test_dir=base_test_env["test_dir"],
|
|
device=base_test_env["device"],
|
|
batch_per_device=batch_per_device,
|
|
)
|
|
|
|
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"""
|
|
grad_accum_steps_list = [1, 2, 4]
|
|
|
|
for grad_accum_steps in grad_accum_steps_list:
|
|
train_config = train_config_factory(
|
|
model_fn=lambda: base_test_env["model"],
|
|
dataset=random_dataset,
|
|
test_dir=base_test_env["test_dir"],
|
|
device=base_test_env["device"],
|
|
batch_per_device=2,
|
|
grad_accum_steps=grad_accum_steps,
|
|
)
|
|
|
|
trainer = Trainer(train_config)
|
|
trainer.train()
|
|
|
|
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_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:
|
|
train_config = train_config_factory(
|
|
model_fn=lambda: base_test_env["model"],
|
|
dataset=random_dataset,
|
|
test_dir=base_test_env["test_dir"],
|
|
device=base_test_env["device"],
|
|
batch_per_device=config["batch_per_device"],
|
|
grad_accum_steps=config["grad_accum_steps"],
|
|
)
|
|
|
|
assert train_config.grad_accum_steps == config["grad_accum_steps"]
|
|
assert train_config.batch_per_device == config["batch_per_device"]
|