fix(finetune): all_reduce val loss across devices for initial and final validation

litgpt/finetune/adapter.py

@@ -226,6 +226,7 @@ def main(
     # Final evaluation
     if eval.final_validation:
         val_loss = validate(fabric, model, val_dataloader, dataclasses.replace(eval, max_iters=len(val_dataloader)))
+        val_loss = fabric.all_reduce(val_loss.detach().clone().to(fabric.device), reduce_op="mean")
         metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
         fabric.log_dict(metrics)
         fabric.print(f"Final evaluation | val loss: {val_loss.item():.3f} | val ppl: {math.exp(val_loss):.3f}")
@@ -268,6 +269,7 @@ def fit(
 
     if eval.initial_validation:
         val_loss = validate(fabric, model, val_dataloader, dataclasses.replace(eval, max_iters=len(val_dataloader)))
+        val_loss = fabric.all_reduce(val_loss.detach().clone().to(fabric.device), reduce_op="mean")
         val_loss = f"{val_loss:.3f}"
     else:
         fabric.print("Verifying settings ...")
@@ -356,16 +358,16 @@ def fit(
             generate_example(fabric, model, tokenizer, eval, data)
             t1 = time.perf_counter() - t0
 
-            val_loss_tensor = val_loss.detach().clone().to(fabric.device)
             val_time_tensor = torch.tensor(t1, device=fabric.device, dtype=torch.float32)
 
-            fabric.all_reduce(val_loss_tensor, reduce_op="mean")
-            fabric.all_reduce(val_time_tensor, reduce_op="mean")
+            # reassign so that the training progress lines also report the reduced loss
+            val_loss = fabric.all_reduce(val_loss.detach().clone().to(fabric.device), reduce_op="mean")
+            val_time_tensor = fabric.all_reduce(val_time_tensor, reduce_op="mean")
 
             fabric.print(
-                f"iter {iter_num}: val loss {val_loss_tensor.item():.4f}, val time: {val_time_tensor.item() * 1000:.2f} ms"
+                f"iter {iter_num}: val loss {val_loss.item():.4f}, val time: {val_time_tensor.item() * 1000:.2f} ms"
             )
-            metrics = {"val_loss": val_loss_tensor, "val_ppl": math.exp(val_loss_tensor)}
+            metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
             fabric.log_dict(metrics, step=iter_num)
             fabric.barrier()
 

litgpt/finetune/adapter_v2.py

@@ -243,6 +243,7 @@ def main(
     # Final evaluation
     if eval.final_validation:
         val_loss = validate(fabric, model, val_dataloader, dataclasses.replace(eval, max_iters=len(val_dataloader)))
+        val_loss = fabric.all_reduce(val_loss.detach().clone().to(fabric.device), reduce_op="mean")
         metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
         fabric.log_dict(metrics)
         fabric.print(f"Final evaluation | val loss: {val_loss.item():.3f} | val ppl: {math.exp(val_loss):.3f}")
@@ -286,6 +287,7 @@ def fit(
 
     if eval.initial_validation:
         val_loss = validate(fabric, model, val_dataloader, dataclasses.replace(eval, max_iters=len(val_dataloader)))
+        val_loss = fabric.all_reduce(val_loss.detach().clone().to(fabric.device), reduce_op="mean")
         val_loss = f"{val_loss:.3f}"
     else:
         fabric.print("Verifying settings ...")
@@ -383,16 +385,16 @@ def fit(
             generate_example(fabric, model, tokenizer, eval, data)
             t1 = time.perf_counter() - t0
 
-            val_loss_tensor = val_loss.detach().clone().to(fabric.device)
             val_time_tensor = torch.tensor(t1, device=fabric.device, dtype=torch.float32)
 
-            fabric.all_reduce(val_loss_tensor, reduce_op="mean")
-            fabric.all_reduce(val_time_tensor, reduce_op="mean")
+            # reassign so that the training progress lines also report the reduced loss
+            val_loss = fabric.all_reduce(val_loss.detach().clone().to(fabric.device), reduce_op="mean")
+            val_time_tensor = fabric.all_reduce(val_time_tensor, reduce_op="mean")
 
             fabric.print(
-                f"iter {iter_num}: val loss {val_loss_tensor.item():.4f}, val time: {val_time_tensor.item() * 1000:.2f} ms"
+                f"iter {iter_num}: val loss {val_loss.item():.4f}, val time: {val_time_tensor.item() * 1000:.2f} ms"
             )
-            metrics = {"val_loss": val_loss_tensor, "val_ppl": math.exp(val_loss_tensor)}
+            metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
             fabric.log_dict(metrics, step=iter_num)
             fabric.barrier()
 

litgpt/finetune/full.py

@@ -191,6 +191,7 @@ def main(
     # Final evaluation
     if eval.final_validation:
         val_loss = validate(fabric, model, val_dataloader, dataclasses.replace(eval, max_iters=len(val_dataloader)))
+        val_loss = fabric.all_reduce(val_loss.detach().clone().to(fabric.device), reduce_op="mean")
         metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
         fabric.log_dict(metrics, step=state["iter_num"])
         fabric.print(f"Final evaluation | val loss: {val_loss.item():.3f} | val ppl: {math.exp(val_loss):.3f}")
@@ -241,6 +242,7 @@ def fit(
 
     if eval.initial_validation:
         val_loss = validate(fabric, model, val_dataloader, dataclasses.replace(eval, max_iters=len(val_dataloader)))
+        val_loss = fabric.all_reduce(val_loss.detach().clone().to(fabric.device), reduce_op="mean")
         val_loss = f"{val_loss:.3f}"
     else:
         fabric.print("Verifying settings ...")
@@ -328,16 +330,16 @@ def fit(
             generate_example(fabric, model, tokenizer, eval, data)
             t1 = time.perf_counter() - t0
 
-            val_loss_tensor = val_loss.detach().clone().to(fabric.device)
             val_time_tensor = torch.tensor(t1, device=fabric.device, dtype=torch.float32)
 
-            fabric.all_reduce(val_loss_tensor, reduce_op="mean")
-            fabric.all_reduce(val_time_tensor, reduce_op="mean")
+            # reassign so that the training progress lines also report the reduced loss
+            val_loss = fabric.all_reduce(val_loss.detach().clone().to(fabric.device), reduce_op="mean")
+            val_time_tensor = fabric.all_reduce(val_time_tensor, reduce_op="mean")
 
             fabric.print(
-                f"iter {state['iter_num']}: val loss {val_loss_tensor.item():.4f}, val time: {val_time_tensor.item() * 1000:.2f} ms"
+                f"iter {state['iter_num']}: val loss {val_loss.item():.4f}, val time: {val_time_tensor.item() * 1000:.2f} ms"
             )
-            metrics = {"val_loss": val_loss_tensor, "val_ppl": math.exp(val_loss_tensor)}
+            metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
             fabric.log_dict(metrics, step=state["iter_num"])
             fabric.barrier()
         if train.save_interval is not None and not is_accumulating and state["step_count"] % train.save_interval == 0:

litgpt/finetune/lora.py

@@ -259,6 +259,7 @@ def main(
     # Final evaluation
     if eval.final_validation:
         val_loss = validate(fabric, model, val_dataloader, dataclasses.replace(eval, max_iters=len(val_dataloader)))
+        val_loss = fabric.all_reduce(val_loss.detach().clone().to(fabric.device), reduce_op="mean")
         metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
         fabric.log_dict(metrics)
         fabric.print(f"Final evaluation | val loss: {val_loss.item():.3f} | val ppl: {math.exp(val_loss):.3f}")
@@ -309,6 +310,7 @@ def fit(
 
     if eval.initial_validation:
         val_loss = validate(fabric, model, val_dataloader, dataclasses.replace(eval, max_iters=len(val_dataloader)))
+        val_loss = fabric.all_reduce(val_loss.detach().clone().to(fabric.device), reduce_op="mean")
         val_loss = f"{val_loss:.3f}"
     else:
         fabric.print("Verifying settings ...")
@@ -409,16 +411,16 @@ def fit(
             fabric.barrier()
             t1 = time.perf_counter() - t0
 
-            val_loss_tensor = val_loss.detach().clone().to(fabric.device)
             val_time_tensor = torch.tensor(t1, device=fabric.device, dtype=torch.float32)
 
-            fabric.all_reduce(val_loss_tensor, reduce_op="mean")
-            fabric.all_reduce(val_time_tensor, reduce_op="mean")
+            # reassign so that the training progress lines also report the reduced loss
+            val_loss = fabric.all_reduce(val_loss.detach().clone().to(fabric.device), reduce_op="mean")
+            val_time_tensor = fabric.all_reduce(val_time_tensor, reduce_op="mean")
 
             fabric.print(
-                f"iter {iter_num}: val loss {val_loss_tensor.item():.4f}, val time: {val_time_tensor.item() * 1000:.2f} ms"
+                f"iter {iter_num}: val loss {val_loss.item():.4f}, val time: {val_time_tensor.item() * 1000:.2f} ms"
             )
-            metrics = {"val_loss": val_loss_tensor, "val_ppl": math.exp(val_loss_tensor)}
+            metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
             fabric.log_dict(metrics, step=iter_num)
             fabric.barrier()
 

tests/test_adapter.py

@@ -1,5 +1,6 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 import os
+import re
 from contextlib import redirect_stdout
 from copy import deepcopy
 from dataclasses import asdict
@@ -107,6 +108,51 @@ 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_all_reduces_val_loss(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)
+
+    # simulate a multi-device run by making the reduced loss differ from the rank-local one
+    def all_reduce_mock(self, data, group=None, reduce_op="mean"):
+        if reduce_op == "mean" and isinstance(data, torch.Tensor):
+            return data + 100
+        return data
+
+    monkeypatch.setattr(Fabric, "all_reduce", all_reduce_mock)
+
+    out_dir = tmp_path / "out"
+    stdout = StringIO()
+    with redirect_stdout(stdout), mock.patch("sys.argv", ["adapter.py", str(fake_checkpoint_dir)]):
+        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=out_dir,
+            precision="32-true",
+            train=TrainArgs(global_batch_size=1, epochs=1, max_steps=4, micro_batch_size=1),
+            eval=EvalArgs(interval=2, max_iters=2, max_new_tokens=1, initial_validation=True),
+        )
+
+    logs = stdout.getvalue()
+    # the initial validation and the training progress lines must report the reduced loss
+    progress_vals = re.findall(r"val: (\d+\.\d+)", logs)
+    assert progress_vals and all(float(v) > 100 for v in progress_vals)
+    # the periodic validation must report the reduced loss
+    periodic_vals = re.findall(r"val loss (\d+\.\d+)", logs)
+    assert periodic_vals and all(float(v) > 100 for v in periodic_vals)
+    # the final validation must report the reduced loss
+    final_vals = re.findall(r"Final evaluation \| val loss: (\d+\.\d+)", logs)
+    assert len(final_vals) == 1 and float(final_vals[0]) > 100
+
+
 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

@@ -1,5 +1,6 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 import os
+import re
 from contextlib import redirect_stdout
 from copy import deepcopy
 from io import StringIO
@@ -124,6 +125,51 @@ 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_all_reduces_val_loss(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)
+
+    # simulate a multi-device run by making the reduced loss differ from the rank-local one
+    def all_reduce_mock(self, data, group=None, reduce_op="mean"):
+        if reduce_op == "mean" and isinstance(data, torch.Tensor):
+            return data + 100
+        return data
+
+    monkeypatch.setattr(Fabric, "all_reduce", all_reduce_mock)
+
+    out_dir = tmp_path / "out"
+    stdout = StringIO()
+    with redirect_stdout(stdout), mock.patch("sys.argv", ["adapter_v2.py", str(fake_checkpoint_dir)]):
+        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=out_dir,
+            precision="32-true",
+            train=TrainArgs(global_batch_size=1, epochs=1, max_steps=4, micro_batch_size=1),
+            eval=EvalArgs(interval=2, max_iters=2, max_new_tokens=1, initial_validation=True),
+        )
+
+    logs = stdout.getvalue()
+    # the initial validation and the training progress lines must report the reduced loss
+    progress_vals = re.findall(r"val: (\d+\.\d+)", logs)
+    assert progress_vals and all(float(v) > 100 for v in progress_vals)
+    # the periodic validation must report the reduced loss
+    periodic_vals = re.findall(r"val loss (\d+\.\d+)", logs)
+    assert periodic_vals and all(float(v) > 100 for v in periodic_vals)
+    # the final validation must report the reduced loss
+    final_vals = re.findall(r"Final evaluation \| val loss: (\d+\.\d+)", logs)
+    assert len(final_vals) == 1 and float(final_vals[0]) > 100
+
+
 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

@@ -1,13 +1,15 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
 import os
+import re
 from contextlib import redirect_stdout
 from io import StringIO
 from unittest import mock
 from unittest.mock import Mock
 
 import torch
 import yaml
+from lightning import Fabric
 
 import litgpt.finetune.full as module
 from litgpt.args import EvalArgs, TrainArgs
@@ -69,3 +71,48 @@ 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_all_reduces_val_loss(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)
+
+    # simulate a multi-device run by making the reduced loss differ from the rank-local one
+    def all_reduce_mock(self, data, group=None, reduce_op="mean"):
+        if reduce_op == "mean" and isinstance(data, torch.Tensor):
+            return data + 100
+        return data
+
+    monkeypatch.setattr(Fabric, "all_reduce", all_reduce_mock)
+
+    out_dir = tmp_path / "out"
+    stdout = StringIO()
+    with redirect_stdout(stdout), mock.patch("sys.argv", ["full.py", str(fake_checkpoint_dir)]):
+        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=out_dir,
+            precision="32-true",
+            train=TrainArgs(global_batch_size=1, epochs=1, max_steps=4, micro_batch_size=1),
+            eval=EvalArgs(interval=2, max_iters=2, max_new_tokens=1, initial_validation=True),
+        )
+
+    logs = stdout.getvalue()
+    # the initial validation and the training progress lines must report the reduced loss
+    progress_vals = re.findall(r"val: (\d+\.\d+)", logs)
+    assert progress_vals and all(float(v) > 100 for v in progress_vals)
+    # the periodic validation must report the reduced loss
+    periodic_vals = re.findall(r"val loss (\d+\.\d+)", logs)
+    assert periodic_vals and all(float(v) > 100 for v in periodic_vals)
+    # the final validation must report the reduced loss
+    final_vals = re.findall(r"Final evaluation \| val loss: (\d+\.\d+)", logs)
+    assert len(final_vals) == 1 and float(final_vals[0]) > 100

tests/test_lora.py

@@ -1,5 +1,6 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 import os
+import re
 from contextlib import redirect_stdout
 from copy import deepcopy
 from io import StringIO
@@ -290,6 +291,52 @@ 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_all_reduces_val_loss(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)
+
+    # simulate a multi-device run by making the reduced loss differ from the rank-local one
+    def all_reduce_mock(self, data, group=None, reduce_op="mean"):
+        if reduce_op == "mean" and isinstance(data, torch.Tensor):
+            return data + 100
+        return data
+
+    monkeypatch.setattr(Fabric, "all_reduce", all_reduce_mock)
+
+    out_dir = tmp_path / "out"
+    stdout = StringIO()
+    with redirect_stdout(stdout), mock.patch("sys.argv", ["lora.py", str(fake_checkpoint_dir)]):
+        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=out_dir,
+            precision="32-true",
+            train=TrainArgs(global_batch_size=1, epochs=1, max_steps=4, micro_batch_size=1),
+            eval=EvalArgs(interval=2, max_iters=2, max_new_tokens=1, initial_validation=True),
+        )
+
+    logs = stdout.getvalue()
+    # the initial validation and the training progress lines must report the reduced loss
+    progress_vals = re.findall(r"val: (\d+\.\d+)", logs)
+    assert progress_vals and all(float(v) > 100 for v in progress_vals)
+    # the periodic validation must report the reduced loss
+    periodic_vals = re.findall(r"val loss (\d+\.\d+)", logs)
+    assert periodic_vals and all(float(v) > 100 for v in periodic_vals)
+    # the final validation must report the reduced loss
+    final_vals = re.findall(r"Final evaluation \| val loss: (\d+\.\d+)", logs)
+    assert len(final_vals) == 1 and float(final_vals[0]) > 100
+
+
 def test_lora_init_when_linear_overridden():
     class MyLinear(torch.nn.Linear):
         def __init__(self, *args, **kwargs):