aggregate val_loss (#1971)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: ysjprojects

litgpt/finetune/adapter.py

@@ -347,8 +347,17 @@ def fit(
             val_loss = validate(fabric, model, val_dataloader, eval)
             generate_example(fabric, model, tokenizer, eval, data)
             t1 = time.perf_counter() - t0
-            fabric.print(f"iter {iter_num}: val loss {val_loss.item():.4f}, val time: {t1 * 1000:.2f} ms")
-            metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
+
+            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")
+
+            fabric.print(
+                f"iter {iter_num}: val loss {val_loss_tensor.item():.4f}, val time: {val_time_tensor.item() * 1000:.2f} ms"
+            )
+            metrics = {"val_loss": val_loss_tensor, "val_ppl": math.exp(val_loss_tensor)}
             fabric.log_dict(metrics, step=iter_num)
             fabric.barrier()
 

litgpt/finetune/adapter_v2.py

@@ -348,8 +348,17 @@ def fit(
             val_loss = validate(fabric, model, val_dataloader, eval)
             generate_example(fabric, model, tokenizer, eval, data)
             t1 = time.perf_counter() - t0
-            fabric.print(f"iter {iter_num}: val loss {val_loss.item():.4f}, val time: {t1 * 1000:.2f} ms")
-            metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
+
+            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")
+
+            fabric.print(
+                f"iter {iter_num}: val loss {val_loss_tensor.item():.4f}, val time: {val_time_tensor.item() * 1000:.2f} ms"
+            )
+            metrics = {"val_loss": val_loss_tensor, "val_ppl": math.exp(val_loss_tensor)}
             fabric.log_dict(metrics, step=iter_num)
             fabric.barrier()
 

litgpt/finetune/full.py

@@ -320,8 +320,17 @@ def fit(
             val_loss = validate(fabric, model, val_dataloader, eval)
             generate_example(fabric, model, tokenizer, eval, data)
             t1 = time.perf_counter() - t0
-            fabric.print(f"iter {state['iter_num']}: val loss {val_loss.item():.4f}, val time: {t1 * 1000:.2f} ms")
-            metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
+
+            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")
+
+            fabric.print(
+                f"iter {state['iter_num']}: val loss {val_loss_tensor.item():.4f}, val time: {val_time_tensor.item() * 1000:.2f} ms"
+            )
+            metrics = {"val_loss": val_loss_tensor, "val_ppl": math.exp(val_loss_tensor)}
             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

@@ -379,8 +379,17 @@ def fit(
             val_loss = validate(fabric, model, val_dataloader, eval)
             generate_example(fabric, model, tokenizer, eval, data)
             t1 = time.perf_counter() - t0
-            fabric.print(f"iter {iter_num}: val loss {val_loss.item():.4f}, val time: {t1 * 1000:.2f} ms")
-            metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
+
+            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")
+
+            fabric.print(
+                f"iter {iter_num}: val loss {val_loss_tensor.item():.4f}, val time: {val_time_tensor.item() * 1000:.2f} ms"
+            )
+            metrics = {"val_loss": val_loss_tensor, "val_ppl": math.exp(val_loss_tensor)}
             fabric.log_dict(metrics, step=iter_num)
             fabric.barrier()