From 202557f187d8c70cefbc0e54554c9e0115344b5e Mon Sep 17 00:00:00 2001 From: mhh111 Date: Sat, 11 Jul 2026 17:44:21 +0800 Subject: [PATCH 1/5] [v1] fix grad_norm reporting and clipping under FSDP2 sharded params clip_grad_norm_ on FSDP2 sharded DTensor parameters returns a per-rank local shard norm (== global_norm / sqrt(shard_size)) because _get_total_norm does not all-reduce squared norms across the fsdp mesh. .item() therefore reads only the local partial, so reported grad_norm scales as 1/sqrt(dp_size) (e.g. 8xdp mbs1 vs 4xdp mbs2 differ by sqrt(2)), and the clip coefficient is applied per-shard -- corrupting the update once max_grad_norm is enabled. Compute the true global norm via get_total_norm + full_tensor (which all-reduces across the fsdp shard mesh), then clip with that scalar via clip_grads_with_norm_. The previous cp_size>1 cross-CP all-reduce is dropped: under the default mp_shard_size=world_size the fsdp mesh already spans CP, so it would over-count grad_norm by sqrt(cp_size) and skew the clip coefficient. --- src/llamafactory/v1/core/base_trainer.py | 34 +++++++++++++++++++----- 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index eda95c7693..9731f0bb37 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/src/llamafactory/v1/core/base_trainer.py @@ -279,12 +279,34 @@ def fit(self) -> None: # deepspeed: engine.step() already ran inside backward at the sync boundary grad_norm = self._deepspeed_engine.get_grad_norm() else: - grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.max_grad_norm).item() - - if self.args.dist_config and self.args.dist_config.get("cp_size", 1) > 1: - grad_norm = grad_norm**2 - grad_norm = DistributedInterface().all_reduce(grad_norm, op=ReduceOp.SUM, dim=Dim.CP) - grad_norm = grad_norm**0.5 + # FSDP2 shards params/grads as DTensors across the fsdp mesh, so each rank only + # holds 1/shard_size of every parameter. `get_total_norm`/`clip_grad_norm_` + # return a *per-rank local shard* norm and `.item()` reads only that local + # partial (== global_norm / sqrt(shard_size)). Using it directly makes the + # reported grad_norm scale as 1/sqrt(dp_size) (e.g. 8xdp mbs1 vs 4xdp mbs2 + # differ by sqrt(2)), and -- worse -- makes `clip_grad_norm_` apply the clip + # coefficient per-shard, corrupting the update once clipping is actually on. + # Fix: reduce to the true global norm first (full_tensor all-reduces across the + # fsdp shard mesh), then clip with that scalar via clip_grads_with_norm_. + from torch.distributed.tensor import DTensor + + grads = [p.grad for p in self.model.parameters() if p.grad is not None] + total_norm = torch.nn.utils.get_total_norm(grads) + if isinstance(total_norm, DTensor): + # full_tensor already all-reduces across the whole fsdp shard mesh. With the + # default mp_shard_size = world_size this mesh spans CP too, so FSDP's + # reduce-scatter has already summed grads across CP -- no separate CP + # reduction is wanted (it would over-count grad_norm by sqrt(cp_size) and + # also skew the clip coefficient below). CP-on and CP-off both report the + # true global norm this way. + total_norm = total_norm.full_tensor() + # pass the (replicated) global norm as a Tensor -- clip_grads_with_norm_ does + # torch.clamp(max_norm / (total_norm + 1e-6), max=1.0), which rejects a bare + # python float. .item() for reporting only, after clipping. + torch.nn.utils.clip_grads_with_norm_( + self.model.parameters(), self.args.max_grad_norm, total_norm + ) + grad_norm = total_norm.item() if not torch.isfinite(torch.tensor(grad_norm)): # type: ignore # pyright: ignore [reportUnknownReturnType] logger.warning_rank0(f"Gradient norm is not finite: {grad_norm}") From af9b7ce5e97517997a45809a1febcbe09342d4a4 Mon Sep 17 00:00:00 2001 From: mhh111 Date: Sat, 11 Jul 2026 17:44:36 +0800 Subject: [PATCH 2/5] [v1] fix learning_rate log truncation for small values (<1e-4) learning_rate was formatted with :.4f, so values below 1e-4 (e.g. 1e-5, 2e-5) printed as "0.0000". Use :.4g for learning_rate so 1e-5 shows as "1e-05" while 1e-4 still shows "0.0001". Other floats (loss, grad_norm) keep :.4f. --- .../v1/utils/callbacks/logging_callback.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/llamafactory/v1/utils/callbacks/logging_callback.py b/src/llamafactory/v1/utils/callbacks/logging_callback.py index f1674cfbb1..6c6278f09f 100644 --- a/src/llamafactory/v1/utils/callbacks/logging_callback.py +++ b/src/llamafactory/v1/utils/callbacks/logging_callback.py @@ -54,7 +54,17 @@ def on_log( # Human-readable output to stdout display_logs = {**logs, "step": state.global_step, "total_steps": state.num_training_steps} - parts = ", ".join(f"{k}: {v:.4f}" if isinstance(v, float) else f"{k}: {v}" for k, v in display_logs.items()) + + def _fmt(k: str, v) -> str: + if not isinstance(v, float): + return f"{k}: {v}" + # learning_rate is often < 1e-4 (e.g. 1e-5); :.4f would print "0.0000". + # Use :.4g so small values show as "1e-05" while 1e-4 still shows "0.0001". + if k == "learning_rate": + return f"{k}: {v:.4g}" + return f"{k}: {v:.4f}" + + parts = ", ".join(_fmt(k, v) for k, v in display_logs.items()) logger.info_rank0(parts) # Append to JSONL log file in output_dir From 6b7bb2b9a27e8b22e873df0474c69cc272993005 Mon Sep 17 00:00:00 2001 From: mhh111 Date: Sat, 11 Jul 2026 17:45:53 +0800 Subject: [PATCH 3/5] [v1] fix num_key_value_heads typo in ulysses sequence parallel setup In apply_sequence_parallel, the try branch assigned num_key_value_heads from model.config.num_attention_heads (typo -- both read num_attention_heads), so the cp_size divisibility assert for KV heads validated the wrong value. The except branch (text_config path) was already correct. --- .../model_plugins/parallelization/sequence_parallel.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py b/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py index 8d5073bb48..095379ae40 100644 --- a/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py +++ b/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py @@ -90,7 +90,10 @@ def apply_sequence_parallel(model, model_args): set_ulysses_sequence_parallel_group(DistributedInterface().get_group(Dim.CP)) try: - num_attention_heads, num_key_value_heads = model.config.num_attention_heads, model.config.num_attention_heads + num_attention_heads, num_key_value_heads = ( + model.config.num_attention_heads, + model.config.num_key_value_heads, + ) except AttributeError: num_attention_heads, num_key_value_heads = ( model.config.text_config.num_attention_heads, From a3c8a5aca16111efd997229c2e6f28eb2cf5051e Mon Sep 17 00:00:00 2001 From: mhh111 Date: Mon, 13 Jul 2026 09:37:57 +0800 Subject: [PATCH 4/5] [v1] shorten grad_norm fix comment --- src/llamafactory/v1/core/base_trainer.py | 25 +++++++----------------- 1 file changed, 7 insertions(+), 18 deletions(-) diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index 9731f0bb37..35972c845c 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/src/llamafactory/v1/core/base_trainer.py @@ -279,30 +279,19 @@ def fit(self) -> None: # deepspeed: engine.step() already ran inside backward at the sync boundary grad_norm = self._deepspeed_engine.get_grad_norm() else: - # FSDP2 shards params/grads as DTensors across the fsdp mesh, so each rank only - # holds 1/shard_size of every parameter. `get_total_norm`/`clip_grad_norm_` - # return a *per-rank local shard* norm and `.item()` reads only that local - # partial (== global_norm / sqrt(shard_size)). Using it directly makes the - # reported grad_norm scale as 1/sqrt(dp_size) (e.g. 8xdp mbs1 vs 4xdp mbs2 - # differ by sqrt(2)), and -- worse -- makes `clip_grad_norm_` apply the clip - # coefficient per-shard, corrupting the update once clipping is actually on. - # Fix: reduce to the true global norm first (full_tensor all-reduces across the - # fsdp shard mesh), then clip with that scalar via clip_grads_with_norm_. + # FSDP2 shards params/grads across the fsdp mesh, so clip_grad_norm_ returns a + # per-rank local shard norm (global / sqrt(shard_size)): reported grad_norm then + # scales as 1/sqrt(dp_size) and the clip coefficient is applied per-shard. Reduce + # to the true global norm first, then clip with it. from torch.distributed.tensor import DTensor grads = [p.grad for p in self.model.parameters() if p.grad is not None] total_norm = torch.nn.utils.get_total_norm(grads) if isinstance(total_norm, DTensor): - # full_tensor already all-reduces across the whole fsdp shard mesh. With the - # default mp_shard_size = world_size this mesh spans CP too, so FSDP's - # reduce-scatter has already summed grads across CP -- no separate CP - # reduction is wanted (it would over-count grad_norm by sqrt(cp_size) and - # also skew the clip coefficient below). CP-on and CP-off both report the - # true global norm this way. + # full_tensor all-reduces across the fsdp mesh (spans CP under default + # mp_shard=world); a separate CP reduce would over-count by sqrt(cp_size). total_norm = total_norm.full_tensor() - # pass the (replicated) global norm as a Tensor -- clip_grads_with_norm_ does - # torch.clamp(max_norm / (total_norm + 1e-6), max=1.0), which rejects a bare - # python float. .item() for reporting only, after clipping. + # pass a Tensor: clip_grads_with_norm_ clamps max_norm / (total_norm + 1e-6). torch.nn.utils.clip_grads_with_norm_( self.model.parameters(), self.args.max_grad_norm, total_norm ) From 936f20ff34a3f7227960cd3f3ac1e7889075da5f Mon Sep 17 00:00:00 2001 From: mhh111 Date: Mon, 13 Jul 2026 09:53:27 +0800 Subject: [PATCH 5/5] [v1] move DTensor import to top-level --- src/llamafactory/v1/core/base_trainer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index 35972c845c..85bca0b8ed 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/src/llamafactory/v1/core/base_trainer.py @@ -31,6 +31,7 @@ import torch import torch.nn.functional as F +from torch.distributed.tensor import DTensor from ..accelerator.helper import ReduceOp from ..accelerator.interface import Dim, DistributedInterface @@ -283,8 +284,6 @@ def fit(self) -> None: # per-rank local shard norm (global / sqrt(shard_size)): reported grad_norm then # scales as 1/sqrt(dp_size) and the clip coefficient is applied per-shard. Reduce # to the true global norm first, then clip with it. - from torch.distributed.tensor import DTensor - grads = [p.grad for p in self.model.parameters() if p.grad is not None] total_norm = torch.nn.utils.get_total_norm(grads) if isinstance(total_norm, DTensor):