fix(finetune): apply train.max_norm gradient clipping in finetune scripts

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:

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:

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:

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:

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)

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)

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)

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):