From e9a04a13f70e9e743f10690f99b5c45c00f226fe Mon Sep 17 00:00:00 2001 From: Philipp Sinitsin Date: Thu, 11 Jun 2026 22:30:04 +0100 Subject: [PATCH] fix(finetune): apply train.max_norm gradient clipping in finetune scripts The four finetune scripts (lora, full, adapter, adapter_v2) listed train.max_norm as unsupported in validate_args even though the finetune configs in config_hub expose it. Remove it from the unsupported lists and clip gradients at optimizer-step boundaries with fabric.clip_gradients, mirroring pretrain.py. The argument remains optional, so behavior is unchanged when it is not set. Add regression tests asserting that clip_gradients is called once per optimizer step with the configured value. --- litgpt/finetune/adapter.py | 4 +++- litgpt/finetune/adapter_v2.py | 4 +++- litgpt/finetune/full.py | 4 +++- litgpt/finetune/lora.py | 4 +++- tests/test_adapter.py | 36 +++++++++++++++++++++++++++++++++++ tests/test_adapter_v2.py | 36 +++++++++++++++++++++++++++++++++++ tests/test_full.py | 36 +++++++++++++++++++++++++++++++++++ tests/test_lora.py | 36 +++++++++++++++++++++++++++++++++++ 8 files changed, 156 insertions(+), 4 deletions(-) diff --git a/litgpt/finetune/adapter.py b/litgpt/finetune/adapter.py index 87ef3c52db..fe62a4a5bc 100644 --- a/litgpt/finetune/adapter.py +++ b/litgpt/finetune/adapter.py @@ -310,6 +310,8 @@ def fit( running_loss.update(loss.detach()) if not is_accumulating: + if train.max_norm is not None: + fabric.clip_gradients(model, optimizer, max_norm=train.max_norm) optimizer.step() optimizer.zero_grad() scheduler.step() @@ -476,7 +478,7 @@ def save_adapter_checkpoint(fabric: L.Fabric, model: torch.nn.Module, file_path: def validate_args(train: TrainArgs, eval: EvalArgs) -> None: issues = [] - unsupported = [(train, ["max_tokens", "max_norm", "tie_embeddings", "lr_warmup_fraction"])] + unsupported = [(train, ["max_tokens", "tie_embeddings", "lr_warmup_fraction"])] for args, names in unsupported: for name in names: if getattr(args, name) is not None: diff --git a/litgpt/finetune/adapter_v2.py b/litgpt/finetune/adapter_v2.py index be7f72e376..33365dd0ef 100644 --- a/litgpt/finetune/adapter_v2.py +++ b/litgpt/finetune/adapter_v2.py @@ -337,6 +337,8 @@ def fit( running_loss.update(loss.detach()) if not is_accumulating: + if train.max_norm is not None: + fabric.clip_gradients(model, optimizer, max_norm=train.max_norm) optimizer.step() optimizer.zero_grad() scheduler.step() @@ -499,7 +501,7 @@ def save_adapter_v2_checkpoint(fabric: L.Fabric, model: torch.nn.Module, file_pa def validate_args(train: TrainArgs, eval: EvalArgs) -> None: issues = [] - unsupported = [(train, ["max_tokens", "max_norm", "tie_embeddings", "lr_warmup_fraction"])] + unsupported = [(train, ["max_tokens", "tie_embeddings", "lr_warmup_fraction"])] for args, names in unsupported: for name in names: if getattr(args, name) is not None: diff --git a/litgpt/finetune/full.py b/litgpt/finetune/full.py index e13fa3f90a..1113b5db8a 100644 --- a/litgpt/finetune/full.py +++ b/litgpt/finetune/full.py @@ -287,6 +287,8 @@ def fit( running_loss.update(loss.detach()) if not is_accumulating: + if train.max_norm is not None: + fabric.clip_gradients(model, optimizer, max_norm=train.max_norm) optimizer.step() optimizer.zero_grad() scheduler.step() @@ -442,7 +444,7 @@ def get_longest_seq_length(data: list[dict]) -> tuple[int, int]: def validate_args(train: TrainArgs, eval: EvalArgs) -> None: issues = [] - unsupported = [(train, ["max_tokens", "max_norm", "tie_embeddings", "lr_warmup_fraction"])] + unsupported = [(train, ["max_tokens", "tie_embeddings", "lr_warmup_fraction"])] for args, names in unsupported: for name in names: if getattr(args, name) is not None: diff --git a/litgpt/finetune/lora.py b/litgpt/finetune/lora.py index fbecf5a815..35a661742f 100644 --- a/litgpt/finetune/lora.py +++ b/litgpt/finetune/lora.py @@ -361,6 +361,8 @@ def fit( running_loss.update(loss.detach()) if not is_accumulating: + if train.max_norm is not None: + fabric.clip_gradients(model, optimizer, max_norm=train.max_norm) optimizer.step() optimizer.zero_grad() scheduler.step() @@ -562,7 +564,7 @@ def save_lora_checkpoint(fabric: L.Fabric, model: torch.nn.Module, file_path: Pa def validate_args(train: TrainArgs, eval: EvalArgs) -> None: issues = [] - unsupported = [(train, ["max_tokens", "max_norm", "tie_embeddings", "lr_warmup_fraction"])] + unsupported = [(train, ["max_tokens", "tie_embeddings", "lr_warmup_fraction"])] for args, names in unsupported: for name in names: if getattr(args, name) is not None: diff --git a/tests/test_adapter.py b/tests/test_adapter.py index 36b320d23e..6edd6c8f8d 100644 --- a/tests/test_adapter.py +++ b/tests/test_adapter.py @@ -107,6 +107,42 @@ def test_adapter_script(tmp_path, fake_checkpoint_dir, monkeypatch, alpaca_path) assert "of trainable parameters: 168" in logs +@mock.patch.dict(os.environ, {"LT_ACCELERATOR": "cpu"}) +def test_adapter_script_max_norm(tmp_path, fake_checkpoint_dir, monkeypatch, alpaca_path): + model_config = dict(block_size=128, n_layer=2, n_embd=8, n_head=4, padded_vocab_size=8, adapter_start_layer=0) + (fake_checkpoint_dir / "model_config.yaml").write_text(yaml.dump(model_config)) + + monkeypatch.setattr(module, "load_checkpoint", Mock()) + + tokenizer_mock = Mock() + tokenizer_mock.return_value = tokenizer_mock + tokenizer_mock.encode = lambda *_, **__: torch.tensor([3, 2, 1]) + monkeypatch.setattr(module, "Tokenizer", tokenizer_mock) + + stdout = StringIO() + with ( + redirect_stdout(stdout), + mock.patch("sys.argv", ["adapter.py", str(fake_checkpoint_dir)]), + mock.patch.object(Fabric, "clip_gradients") as clip_mock, + ): + module.setup( + fake_checkpoint_dir, + data=Alpaca( + download_dir=alpaca_path.parent, file_name=alpaca_path.name, val_split_fraction=0.5, num_workers=0 + ), + out_dir=tmp_path / "out", + precision="32-true", + train=TrainArgs( + global_batch_size=1, save_interval=2, epochs=1, max_steps=2, micro_batch_size=1, max_norm=1.0 + ), + eval=EvalArgs(interval=2, max_iters=2, max_new_tokens=1), + ) + + # gradient clipping is applied once per optimizer step + assert clip_mock.call_count == 2 + assert all(call.kwargs["max_norm"] == 1.0 for call in clip_mock.call_args_list) + + def test_adapter_gpt_init_weights(): config = Config(n_layer=1, n_head=6, n_embd=12, block_size=1, vocab_size=1, adapter_start_layer=0) model = GPT(config) diff --git a/tests/test_adapter_v2.py b/tests/test_adapter_v2.py index 3a7d17d5e5..4111289fb1 100644 --- a/tests/test_adapter_v2.py +++ b/tests/test_adapter_v2.py @@ -124,6 +124,42 @@ def test_adapter_v2_script(tmp_path, fake_checkpoint_dir, monkeypatch, alpaca_pa assert "of trainable parameters: 552" in logs +@mock.patch.dict(os.environ, {"LT_ACCELERATOR": "cpu"}) +def test_adapter_v2_script_max_norm(tmp_path, fake_checkpoint_dir, monkeypatch, alpaca_path): + model_config = dict(block_size=128, n_layer=2, n_embd=8, n_head=4, padded_vocab_size=8, adapter_start_layer=0) + (fake_checkpoint_dir / "model_config.yaml").write_text(yaml.dump(model_config)) + + monkeypatch.setattr(module, "load_checkpoint", Mock()) + + tokenizer_mock = Mock() + tokenizer_mock.return_value = tokenizer_mock + tokenizer_mock.encode = lambda *_, **__: torch.tensor([3, 2, 1]) + monkeypatch.setattr(module, "Tokenizer", tokenizer_mock) + + stdout = StringIO() + with ( + redirect_stdout(stdout), + mock.patch("sys.argv", ["adapter_v2.py", str(fake_checkpoint_dir)]), + mock.patch.object(Fabric, "clip_gradients") as clip_mock, + ): + module.setup( + fake_checkpoint_dir, + data=Alpaca( + download_dir=alpaca_path.parent, file_name=alpaca_path.name, val_split_fraction=0.5, num_workers=0 + ), + out_dir=tmp_path / "out", + precision="32-true", + train=TrainArgs( + global_batch_size=1, save_interval=2, epochs=1, max_steps=2, micro_batch_size=1, max_norm=1.0 + ), + eval=EvalArgs(interval=2, max_iters=2, max_new_tokens=1), + ) + + # gradient clipping is applied once per optimizer step + assert clip_mock.call_count == 2 + assert all(call.kwargs["max_norm"] == 1.0 for call in clip_mock.call_args_list) + + def test_adapter_v2_gpt_init_weights(): config = Config(n_layer=1, n_head=6, n_embd=12, block_size=1, vocab_size=1, adapter_start_layer=0) model = AdapterV2GPT(config) diff --git a/tests/test_full.py b/tests/test_full.py index 2b0d55a8b9..ee8831613e 100644 --- a/tests/test_full.py +++ b/tests/test_full.py @@ -8,6 +8,7 @@ import torch import yaml +from lightning import Fabric import litgpt.finetune.full as module from litgpt.args import EvalArgs, TrainArgs @@ -69,3 +70,38 @@ def test_full_script(tmp_path, fake_checkpoint_dir, monkeypatch, alpaca_path): assert f"Resuming training from {out_dir / 'step-000006' / 'lit_model.pth'}" in logs assert logs.count("(step)") == 2 assert out_dir / "step-000008" in set(out_dir.iterdir()) + + +@mock.patch.dict(os.environ, {"LT_ACCELERATOR": "cpu"}) +def test_full_script_max_norm(tmp_path, fake_checkpoint_dir, monkeypatch, alpaca_path): + model_config = dict(block_size=128, n_layer=2, n_embd=8, n_head=4, padded_vocab_size=8) + (fake_checkpoint_dir / "model_config.yaml").write_text(yaml.dump(model_config)) + monkeypatch.setattr(module, "load_checkpoint", Mock()) + + tokenizer_mock = Mock() + tokenizer_mock.return_value = tokenizer_mock + tokenizer_mock.encode = lambda *_, **__: torch.tensor([3, 2, 1]) + monkeypatch.setattr(module, "Tokenizer", tokenizer_mock) + + stdout = StringIO() + with ( + redirect_stdout(stdout), + mock.patch("sys.argv", ["full.py", str(fake_checkpoint_dir)]), + mock.patch.object(Fabric, "clip_gradients") as clip_mock, + ): + module.setup( + fake_checkpoint_dir, + data=Alpaca( + download_dir=alpaca_path.parent, file_name=alpaca_path.name, val_split_fraction=0.5, num_workers=0 + ), + out_dir=tmp_path / "out", + precision="32-true", + train=TrainArgs( + global_batch_size=1, save_interval=2, epochs=1, max_steps=2, micro_batch_size=1, max_norm=1.0 + ), + eval=EvalArgs(interval=2, max_iters=2, max_new_tokens=1), + ) + + # gradient clipping is applied once per optimizer step + assert clip_mock.call_count == 2 + assert all(call.kwargs["max_norm"] == 1.0 for call in clip_mock.call_args_list) diff --git a/tests/test_lora.py b/tests/test_lora.py index 46ba75b384..49da157ba8 100644 --- a/tests/test_lora.py +++ b/tests/test_lora.py @@ -290,6 +290,42 @@ def test_lora_script(tmp_path, fake_checkpoint_dir, monkeypatch, alpaca_path): assert "of trainable parameters: 512" in logs +@mock.patch.dict(os.environ, {"LT_ACCELERATOR": "cpu"}) +def test_lora_script_max_norm(tmp_path, fake_checkpoint_dir, monkeypatch, alpaca_path): + model_config = dict(block_size=128, n_layer=2, n_embd=8, n_head=4, padded_vocab_size=8) + (fake_checkpoint_dir / "model_config.yaml").write_text(yaml.dump(model_config)) + monkeypatch.setattr(module, "load_checkpoint", Mock()) + monkeypatch.setattr(module, "merge_lora", Mock()) + + tokenizer_mock = Mock() + tokenizer_mock.return_value = tokenizer_mock + tokenizer_mock.encode = lambda *_, **__: torch.tensor([3, 2, 1]) + monkeypatch.setattr(module, "Tokenizer", tokenizer_mock) + + stdout = StringIO() + with ( + redirect_stdout(stdout), + mock.patch("sys.argv", ["lora.py", str(fake_checkpoint_dir)]), + mock.patch.object(Fabric, "clip_gradients") as clip_mock, + ): + module.setup( + fake_checkpoint_dir, + data=Alpaca( + download_dir=alpaca_path.parent, file_name=alpaca_path.name, val_split_fraction=0.5, num_workers=0 + ), + out_dir=tmp_path / "out", + precision="32-true", + train=TrainArgs( + global_batch_size=1, save_interval=2, epochs=1, max_steps=2, micro_batch_size=1, max_norm=1.0 + ), + eval=EvalArgs(interval=2, max_iters=2, max_new_tokens=1), + ) + + # gradient clipping is applied once per optimizer step + assert clip_mock.call_count == 2 + assert all(call.kwargs["max_norm"] == 1.0 for call in clip_mock.call_args_list) + + def test_lora_init_when_linear_overridden(): class MyLinear(torch.nn.Linear): def __init__(self, *args, **kwargs):