Add ability to perform local gradient accumulation in FP32 for a subset of parameters in the model - #4028

Merged

deepakn94 merged 3 commits into

NVIDIA:mainfrom

deepakn94:dnarayanan/local_grad_accumulation_in_fp32

Mar 25, 2026

Commits on Mar 25, 2026

Create separate higher-precision main_grad tensor for local gradient accumulation
deepakn94
committed
739988f
Add unit tests for FP32 local gradient accumulation feature

deepakn94
and
claude
committed
6fdd1a1
Linting and code review
deepakn94
committed
1df3f83