ViperEkura
5ba21f4eb3
- Add tests/helpers.py with shared config, dataset, tokenizer, executor, and assertion helpers - Replace 15 copies of device one-liner with session-scoped fixture - Collapse 5 near-identical Dataset subclasses into RandomTokenDataset - Remove duplicate _make_config/_make_model/_make_frozen and FakeTokenizer/FakeExecutor definitions - Make test_callbacks and test_early_stopping use existing train_config_factory - Replace 6 duplicate meta.json read blocks with load_shard_meta - Fix mkdtemp leaks in test_lora.py with TemporaryDirectory
140 lines
4.2 KiB
Python
140 lines
4.2 KiB
Python
import torch
|
|
|
|
from astrai.model.components.decoder_block import DecoderBlock
|
|
from astrai.trainer.train_callback import GradientCheckpointingCallback, TrainCallback
|
|
from astrai.trainer.trainer import Trainer
|
|
|
|
|
|
def test_gradient_checkpointing_enable_disable(test_model):
|
|
"""Enable wraps forward, _disable restores it."""
|
|
model = test_model["model"]
|
|
callback = GradientCheckpointingCallback(modules=[DecoderBlock])
|
|
|
|
originals = [layer.forward for layer in model.layers]
|
|
|
|
for layer in model.layers:
|
|
callback._enable(layer)
|
|
|
|
for layer in model.layers:
|
|
assert hasattr(layer, "_original_forward")
|
|
assert layer.forward is not originals[0]
|
|
|
|
for layer in model.layers:
|
|
callback._disable(layer)
|
|
|
|
for layer in model.layers:
|
|
assert not hasattr(layer, "_original_forward")
|
|
|
|
|
|
def test_gradient_checkpointing_empty_modules_noop(test_model):
|
|
"""modules=None should leave forwards untouched."""
|
|
model = test_model["model"]
|
|
callback = GradientCheckpointingCallback()
|
|
|
|
originals = [layer.forward for layer in model.layers]
|
|
|
|
for layer in model.layers:
|
|
callback._enable(layer)
|
|
|
|
for layer, orig in zip(model.layers, originals):
|
|
assert layer.forward is orig
|
|
|
|
|
|
def test_gradient_checkpointing_forward_unchanged(test_model):
|
|
"""Forward output unchanged after patching (no_grad)."""
|
|
model = test_model["model"]
|
|
device = test_model["device"]
|
|
callback = GradientCheckpointingCallback(modules=[DecoderBlock])
|
|
|
|
input_ids = torch.randint(0, 1000, (2, 32)).to(device)
|
|
|
|
with torch.no_grad():
|
|
ref = model(input_ids)["logits"].clone()
|
|
|
|
for layer in model.layers:
|
|
callback._enable(layer)
|
|
|
|
with torch.no_grad():
|
|
out = model(input_ids)["logits"]
|
|
|
|
assert torch.equal(ref, out)
|
|
|
|
|
|
def test_gradient_checkpointing_backward(test_model):
|
|
"""backward passes gradients through checkpointed layers."""
|
|
model = test_model["model"]
|
|
device = test_model["device"]
|
|
callback = GradientCheckpointingCallback(modules=[DecoderBlock])
|
|
|
|
for layer in model.layers:
|
|
callback._enable(layer)
|
|
|
|
input_ids = torch.randint(0, 1000, (2, 32)).to(device)
|
|
target_ids = torch.randint(0, 1000, (2, 32)).to(device)
|
|
|
|
logits = model(input_ids)["logits"]
|
|
loss = torch.nn.functional.cross_entropy(
|
|
logits.flatten(0, 1).float(), target_ids.flatten()
|
|
)
|
|
loss.backward()
|
|
|
|
for name, param in model.named_parameters():
|
|
if param.requires_grad:
|
|
assert param.grad is not None, f"{name} gradient is None"
|
|
|
|
for layer in model.layers:
|
|
callback._disable(layer)
|
|
|
|
model.zero_grad()
|
|
for name, p in model.named_parameters():
|
|
assert p.grad is None or p.grad.sum().item() == 0, f"{name} grad not zeroed"
|
|
|
|
|
|
def test_gradient_checkpointing_trainer_integration(
|
|
base_test_env, random_dataset, train_config_factory, device
|
|
):
|
|
"""Gradient checkpointing runs end-to-end via Trainer."""
|
|
train_config = train_config_factory(
|
|
model_fn=lambda: base_test_env["model"],
|
|
dataset=random_dataset,
|
|
test_dir=base_test_env["test_dir"],
|
|
device=device,
|
|
ckpt_interval=3,
|
|
gradient_checkpointing_modules=[DecoderBlock],
|
|
)
|
|
|
|
trainer = Trainer(train_config)
|
|
trainer.train()
|
|
|
|
|
|
def test_callback_integration(
|
|
base_test_env, random_dataset, train_config_factory, device
|
|
):
|
|
"""Test that all callbacks are properly integrated"""
|
|
train_config = train_config_factory(
|
|
model_fn=lambda: base_test_env["model"],
|
|
dataset=random_dataset,
|
|
test_dir=base_test_env["test_dir"],
|
|
device=device,
|
|
ckpt_interval=3,
|
|
)
|
|
|
|
callback_calls = []
|
|
|
|
class TrackingCallback(TrainCallback):
|
|
def on_train_begin(self, context):
|
|
callback_calls.append("on_train_begin")
|
|
|
|
def on_batch_end(self, context):
|
|
callback_calls.append("on_batch_end")
|
|
|
|
def on_epoch_end(self, context):
|
|
callback_calls.append("on_epoch_end")
|
|
|
|
trainer = Trainer(train_config, callbacks=[TrackingCallback()])
|
|
trainer.train()
|
|
|
|
assert "on_train_begin" in callback_calls
|
|
assert "on_batch_end" in callback_calls
|
|
assert "on_epoch_end" in callback_calls
|