diff --git a/specforge/optimizer.py b/specforge/optimizer.py index 9fc877b7c..f4ba807f9 100644 --- a/specforge/optimizer.py +++ b/specforge/optimizer.py @@ -1,4 +1,5 @@ import torch +import torch.distributed as dist from specforge.lr_scheduler import CosineAnnealingWarmupLR from specforge.utils import print_on_rank0 @@ -36,13 +37,34 @@ def __init__( warmup_steps=int(warmup_ratio * total_steps), ) + def _clip_grad_norm(self): + """Clip by the global grad norm, accumulated across ranks. + + Under FSDP each rank holds only its shard of the gradients, so + `torch.nn.utils.clip_grad_norm_` would compute a rank-local norm and + scale each shard by a different coefficient. + """ + grads = [mp.grad for mp in self.fp32_params if mp.grad is not None] + device = self.fp32_params[0].device if self.fp32_params else None + if grads: + total_norm_sq = torch.stack([g.pow(2).sum() for g in grads]).sum() + else: + total_norm_sq = torch.zeros((), dtype=torch.float32, device=device) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(total_norm_sq, op=dist.ReduceOp.SUM) + total_norm = total_norm_sq.sqrt() + clip_coef = torch.clamp(self.max_grad_norm / (total_norm + 1e-6), max=1.0) + for g in grads: + g.mul_(clip_coef) + return total_norm + def step(self): with torch.no_grad(): for p, mp in zip(self.model_params, self.fp32_params): mp.grad = ( p.grad.detach().to(torch.float32) if p.grad is not None else None ) - grad_norm = torch.nn.utils.clip_grad_norm_(self.fp32_params, self.max_grad_norm) + grad_norm = self._clip_grad_norm() self.last_grad_norm = grad_norm.detach() self.optimizer.step() self.optimizer.zero_grad() diff --git a/specforge/runtime/training/backend.py b/specforge/runtime/training/backend.py index bd314fcf4..d180e3e33 100644 --- a/specforge/runtime/training/backend.py +++ b/specforge/runtime/training/backend.py @@ -214,20 +214,12 @@ def backward(self, loss: torch.Tensor) -> None: loss.backward() def step(self) -> Optional[torch.Tensor]: - """Optimizer step + the distributed grad-norm reduction (run_backward_and_update).""" + """Optimizer step (run_backward_and_update); the optimizer clips by and returns the global grad norm.""" if self.optimizer is None: raise RuntimeError( "FSDPTrainingBackend.step called before optimizer is set" ) - grad_norm = self.optimizer.step() - if grad_norm is not None and dist.is_initialized(): - grad_norm = grad_norm.detach().float() - if torch.cuda.is_available(): - grad_norm = grad_norm.to(torch.cuda.current_device()) - grad_norm = grad_norm.pow(2) - dist.all_reduce(grad_norm, op=dist.ReduceOp.SUM) - grad_norm = grad_norm.sqrt() - return grad_norm + return self.optimizer.step() def state_dict(self) -> dict: if self.module is None: diff --git a/tests/test_optimizer/__init__.py b/tests/test_optimizer/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py b/tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py new file mode 100644 index 000000000..fd3cd2a12 --- /dev/null +++ b/tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py @@ -0,0 +1,91 @@ +import os +import tempfile +import unittest + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from specforge.optimizer import BF16Optimizer + + +def _make_optimizer(seed=0): + torch.manual_seed(seed) + model = torch.nn.Linear(8, 8, bias=False) + return model, BF16Optimizer(model, lr=1e-3, max_grad_norm=0.5) + + +class TestClipGradNormSingleProcess(unittest.TestCase): + def test_matches_torch_reference(self): + """Without distributed, clipping must match torch.nn.utils.clip_grad_norm_.""" + model, opt = _make_optimizer() + grad = torch.randn(8, 8) + + reference = [p.detach().clone().requires_grad_(True) for p in model.parameters()] + for rp in reference: + rp.grad = grad.clone() + ref_norm = torch.nn.utils.clip_grad_norm_(reference, opt.max_grad_norm) + + for mp_ in opt.fp32_params: + mp_.grad = grad.clone() + norm = opt._clip_grad_norm() + + torch.testing.assert_close(norm, ref_norm) + for mp_, rp in zip(opt.fp32_params, reference): + torch.testing.assert_close(mp_.grad, rp.grad) + + def test_step_returns_norm(self): + model, opt = _make_optimizer() + for p in model.parameters(): + p.grad = torch.full_like(p, 0.1) + norm = opt.step() + expected = torch.full((8, 8), 0.1).norm() + torch.testing.assert_close(norm, expected) + + +def _dist_worker(rank, world_size, init_file, results): + dist.init_process_group( + backend="gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + ) + try: + _, opt = _make_optimizer() + # Simulate FSDP sharding: each rank holds a disjoint gradient shard + # with a different magnitude, so local norms differ across ranks. + grad_value = 1.0 if rank == 0 else 2.0 + for mp_ in opt.fp32_params: + mp_.grad = torch.full_like(mp_, grad_value) + norm = opt._clip_grad_norm() + results[rank] = (norm.item(), opt.fp32_params[0].grad.flatten()[0].item()) + finally: + dist.destroy_process_group() + + +class TestClipGradNormDistributed(unittest.TestCase): + def test_global_norm_across_ranks(self): + """Every rank must clip by the same GLOBAL norm, not its local norm.""" + world_size = 2 + with tempfile.TemporaryDirectory() as tmpdir: + init_file = os.path.join(tmpdir, "init") + manager = mp.Manager() + results = manager.dict() + mp.spawn( + _dist_worker, + args=(world_size, init_file, results), + nprocs=world_size, + join=True, + ) + + # 64 elements of 1.0 on rank 0 + 64 elements of 2.0 on rank 1 + global_norm = (64 * 1.0**2 + 64 * 2.0**2) ** 0.5 + clip_coef = 0.5 / (global_norm + 1e-6) + for rank, grad_value in ((0, 1.0), (1, 2.0)): + norm, clipped_first = results[rank] + self.assertAlmostEqual(norm, global_norm, places=4) + self.assertAlmostEqual(clipped_first, grad_value * clip_coef, places=6) + + +if __name__ == "__main__": + unittest.main()