From 75a0864f47e88db0512f16b215e96fdbf277bd5a Mon Sep 17 00:00:00 2001 From: Philipp Sinitsin Date: Sat, 13 Jun 2026 14:47:20 +0100 Subject: [PATCH] fix(finetune): all_reduce val loss across devices for initial and final validation The initial and final validation in the four finetune scripts logged the rank-local loss without reducing it across devices, unlike the periodic validation. The periodic validation also reduced into a temporary tensor that was never assigned back to val_loss, so the training progress lines kept reporting the rank-local value. Reduce the loss in all three places, using the tensor returned by Fabric.all_reduce, and reuse it in the progress lines. All reductions are no-ops on a single device. Adds a CPU regression test per script that stubs Fabric.all_reduce and asserts every printed and logged val loss is the reduced value. --- litgpt/finetune/adapter.py | 12 +++++---- litgpt/finetune/adapter_v2.py | 12 +++++---- litgpt/finetune/full.py | 12 +++++---- litgpt/finetune/lora.py | 12 +++++---- tests/test_adapter.py | 46 ++++++++++++++++++++++++++++++++++ tests/test_adapter_v2.py | 46 ++++++++++++++++++++++++++++++++++ tests/test_full.py | 47 +++++++++++++++++++++++++++++++++++ tests/test_lora.py | 47 +++++++++++++++++++++++++++++++++++ 8 files changed, 214 insertions(+), 20 deletions(-) diff --git a/litgpt/finetune/adapter.py b/litgpt/finetune/adapter.py index 87ef3c52db..2c08826fd5 100644 --- a/litgpt/finetune/adapter.py +++ b/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() diff --git a/litgpt/finetune/adapter_v2.py b/litgpt/finetune/adapter_v2.py index be7f72e376..f261ac4cf0 100644 --- a/litgpt/finetune/adapter_v2.py +++ b/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() diff --git a/litgpt/finetune/full.py b/litgpt/finetune/full.py index e13fa3f90a..290b771398 100644 --- a/litgpt/finetune/full.py +++ b/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: diff --git a/litgpt/finetune/lora.py b/litgpt/finetune/lora.py index fbecf5a815..4688895376 100644 --- a/litgpt/finetune/lora.py +++ b/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() diff --git a/tests/test_adapter.py b/tests/test_adapter.py index 36b320d23e..0ffbacd251 100644 --- a/tests/test_adapter.py +++ b/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) diff --git a/tests/test_adapter_v2.py b/tests/test_adapter_v2.py index 3a7d17d5e5..80de4bb3e6 100644 --- a/tests/test_adapter_v2.py +++ b/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) diff --git a/tests/test_full.py b/tests/test_full.py index 2b0d55a8b9..30927784b8 100644 --- a/tests/test_full.py +++ b/tests/test_full.py @@ -1,6 +1,7 @@ # 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 @@ -8,6 +9,7 @@ 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 diff --git a/tests/test_lora.py b/tests/test_lora.py index 46ba75b384..a153a7d15f 100644 --- a/tests/test_lora.py +++ b/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):