Fix v1 grad norm and lr log

[HelloWorldBeginner]

---

What does this PR do?

Fixes three small bugs in the v1 trainer, all found while debugging cross-layout (different dp_size / cp_size) precision alignment.

1. grad_norm mis-reported (and mis-clipped) under FSDP2 sharded params

File: src/llamafactory/v1/core/base_trainer.py

Under FSDP2, parameters/gradients are sharded DTensors across the fsdp mesh, so each rank holds only 1/shard_size of every parameter. torch.nn.utils.clip_grad_norm_ internally calls _get_total_norm, which computes the norm over the local shard only and does not all-reduce squared norms across the mesh. .item() therefore reads the local partial, which equals global_norm / sqrt(shard_size).

Consequences:

• Reported grad_norm scales as 1/sqGPU mbs1vs 4-GPUmbs2(same global batch,cp_size=1) differ by exactly sqrt(2)(=sqrt(dp8/dp4)), while loss` is identical.
• Worse: the clip coefficient max_norm / total_norm is also computed from the local shard norm, so once max_grad_norm is actually enabled the gradient is clipped per-shard (different coefficient per rank), corrupting the update.

Fix: compute the local norm via get_total_norm, reduce it to the true global norm with .full_tensor() (which all-reduces across the fsdp shard mesh), then clip with that scalar via
clip_grads_with_norm_. .item() is r clipping.

The previous cp_size > 1 cross-CP all-reduce of grad_norm is removed: under the default mp_shard_size = world_size the fsdp mesh already spter has already summed gradientsacross CP — the extra reduce would over-count grad_norm by sqrt(cp_size) and skew the clip coefficient.

Verification: a 2-process DTensor test confirms get_total_norm(...).item() returns the local shard nowhile .full_tensor().item() returns he fix, 8-GPU-mbs1 and 4-GPU-mbs2report the same grad_norm.

2. learning_rate log truncation for small values

File: src/llamafactory/v1/utils/callbacks/logging_callback.py
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) kee

3. num_key_value_heads typo in ulysses sequence-parallel setup

File: src/llamafactory/v1/pluginsn/sequence_parallel.py

In apply_sequence_parallel, the try branch assigned both num_attention_heads and num_key_value_heads from model.conf — the second should benum_key_value_heads). As a result the cp_sizedivisibility assertion for KV heads validated the wrong value. Theexceptbranch (thetext_config` path) was already correct.

Fixes #N/A (no linked issue; found during internal precision debugging)

Before submitting

• Did you read the contributor guideline?
• Did you write any new necessary tests?

[gemini-code-assist]

Code Review

This pull request introduces several improvements and bug fixes: it updates gradient norm calculation and clipping in base_trainer.py to correctly handle FSDP2 DTensor shards; fixes a typo in sequence_parallel.py where num_key_value_heads was incorrectly assigned num_attention_heads; and refactors logging_callback.py to format small learning rates using scientific notation (:.4g) instead of rounding them to zero. Feedback on the pull request points out that torch.nn.utils.get_total_norm is not a public function and will raise an AttributeError, suggesting the use of the private helper _get_total_norm from torch.nn.utils.clip_grad instead.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Calling torch.nn.utils.get_total_norm will raise an AttributeError because get_total_norm is not a public function in torch.nn.utils. The correct private helper is _get_total_norm and it must be imported from torch.nn.utils.clip_grad.

                    from torch.distributed.tensor import DTensor
                    from torch.nn.utils.clip_grad import _get_total_norm

                    grads = [p.grad for p in self.model.parameters() if p.grad is not None]
                    total_norm = _get_total_norm(grads)