Repository: deepspeedai/DeepSpeed
Pull request: Fix: Update grad norm calculation for CPU offload
## Description
This PR fixes an issue where gradient clipping modifications are not reflected in the global gradient norm calculation when CPU offloading is enabled. The issue occurs because the `averaged_gradients` are not being updated with the clipped gradients when CPU offloading is active.

## Problem
When using CPU offloading with gradient clipping:
1. The gradients are successfully clipped using `safe_set_local_grad`
2. However, the `_global_grad_norm` calculation still uses the original unclipped gradients.
3. This leads to incorrect gradient norm reporting and potential issues with gradient clipping effectiveness

## Solution
The fix ensures that the `averaged_gradients` are properly updated with the clipped gradients when CPU offloading is enabled, similar to how it works when CPU offloading is disabled.

## Testing
The fix has been tested with:
- CPU offloading enabled and disabled
- Different gradient clipping values
- A simple model with linear layers
- Both FP16 and BF16

## Related Issues
Fixes #7292 

Problem statement:
Question about `safe_set_local_grad` and `safe_get_local_grad`
deepspeed==0.16.7
I'm using zero stage-3

I write a custom Trainer to do some custom grad clipping after `backward()` is called and before `step()`. I use `safe_get_local_grad` and `safe_set_local_grad` to do the custom grad clipping. After clipping, I call `safe_get_local_grad` again, and the result is as expected. But then i add a callback on the end of `step()` and print the `ds_engine.optimizer._global_grad_norm`, it equals to the original grad(before the clipping). Does it mean I fail to set new value to grad?

Hints from issue discussion:
@breaddaerb, thanks for your question. You are correct in that we don't propagate the update gradient back into `.global_grad_norm`. Are you able to provide a PR for this?
I take a closer look at `DeepSpeedZeroOptimizer_Stage3`. The `_global_grad_norm` is calculated in the `step()` by calling ` _get_norm_groups()` first, and it takes `averaged_gradients` and `fp16_groups` to get norm groups.
https://github.com/deepspeedai/DeepSpeed/blob/b048cc2b46a559b8f81c4315008cec34e650bbe0/deepspeed/runtime/zero/stage3.py#L2095
https://github.com/deepspeedai/DeepSpeed/blob/b048cc2b46a559b8f81c4315008cec34e650bbe0/deepspeed/runtime/zero/stage3.py#L1887
Does it mean the gradient is not changed at all or just the `_global_grad_norm` is not updated with the new grad?
@breaddaerb, since multiple data structures are involved it is hard to say without repro code. Are you able to provide some test code?
Hi @breaddaerb / @sfc-gh-truwase,

I've written a piece of code to repro this.

```
import torch
import torch.nn as nn
import deepspeed
from deepspeed.utils import safe_get_local_grad, safe_set_local_grad
import argparse

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin1 = nn.Linear(4, 8)
        self.lin2 = nn.Linear(8, 2)

    def forward(self, x):
        x = self.lin1(x)
        x = self.lin2(x)
        return x

def get_deepspeed_config():
    return {
        "train_batch_size": 8,
        "gradient_clipping": 1.0,
        "zero_optimization": {
            "stage": 3,
            "offload_optimizer": {"device": "cpu"},
            "overlap_comm": False,
            "contiguous_gradients": True,
            "reduce_bucket_size": 5e7,
            "stage3_prefetch_bucket_size": 5e7,
            "stage3_param_persistence_threshold": 1e4
        },
        "fp16": {"enabled": True, "initial_scale_power": 12},
        "gradient_accumulation_steps": 1,
        "steps_per_print": 1,
        "optimizer": {
            "type": "Adam",
            "params": {
                "lr": 1e-3,
                "betas": [0.9, 0.999],
                "eps": 1e-8,
                "weight_decay": 0
            }
        }
    }

def custom_grad_clip(model_engine, clip_value=0.1):
    modified_count = 0
    
    for param in model_engine.parameters():
        if not hasattr(param, 'ds_id'):
            continue
            
        grad = safe_get_local_grad(param)
        if grad is not None:
            # Apply clipping
            clamped_grad = torch.clamp(grad, -clip_value, clip_value)
            safe_set_local_grad(param, clamped_grad)
            modified_count += 1
    
    return modified_count

def get_grad_norms(model_engine):
    try:
        norm_groups = model_engine.optimizer._get_norm_groups()
        global_norm = torch.linalg.vector_norm(torch.stack(norm_groups))

        return global_norm, norm_groups, model_engine.optimizer._global_grad_norm
    except Exception as e:
        print(f"Warning: Could not get gradient norms: {e}")
        return None, None

def main():
    # Initialize model and DeepSpeed
    model = SimpleModel()
    parser = argparse.ArgumentParser()
    parser.add_argument('--local_rank', type=int, default=-1, help="Local rank for distributed training")
    args, _ = parser.parse_known_args()

    # Initialize DeepSpeed
    model_engine, optimizer, _, _ = deepspeed.initialize(
        args=args,
        model=model,
        config=get_deepspeed_config()
    )

    model_engine.train()

    # Create dummy data
    torch.manual_seed(10)
    inputs = torch.randn(8, 4, device=model_engine.device)
    targets = torch.randn(8, 2, device=model_engine.device)

    if model_engine.fp16_enabled():
        inputs = inputs.half()
        targets = targets.half()

    # Forward pass
    outputs = model_engine(inputs)
    loss = nn.MSELoss()(outputs, targets)
    print(f"Initial loss: {loss.item():.6f}")

    # Backward pass
    model_engine.backward(loss)

    # Get initial global norm
    initial_global_norm = model_engine.optimizer._global_grad_norm
    print(f"Initial global grad norm: {initial_global_norm:.6f}")

    # Get norm groups before clipping
    pre_clip_global_norm, pre_clip_norm_groups, pre_clip_global_grad_norm = get_grad_norms(model_engine)
    if pre_clip_global_norm is not None:
        print(f"Pre-clipping global norm: {pre_clip_global_norm:.6f}")

    # Apply custom gradient clipping
    modified_count = custom_grad_clip(model_engine, clip_value=0.1)
    print(f"Modified {modified_count} parameter gradients")

    # Get norm groups after clipping
    post_clip_global_norm, post_clip_norm_groups, post_clip_global_grad_norm = get_grad_norms(model_engine)
    if post_clip_global_norm is not None:
        print(f"Post-clipping global norm: {post_clip_global_norm:.6f}")

    # Run optimizer step
    model_engine.step()
    print("Optimizer step completed successfully")

    final_global_norm = model_engine.optimizer._global_grad_norm


    # Print summary
    print("\nSummary:")
    if pre_clip_norm_groups is not None:
        print(f"Pre-clipping norms: {[norm.item() for norm in pre_clip_norm_groups]}")
    if post_clip_norm_groups is not None:
        print(f"Post-clipping norms: {[norm.item() for norm in post_clip_norm_groups]}")


    print(f"Pre-clipping global grad norm: {pre_clip_global_grad_norm:.6f}")
    print(f"Post-clipping global grad norm: {post_clip_global_grad_norm:.6f}")
    print(f"Final global grad norm: {final_global_norm:.6f}")


if __name__ == "__main__":
    main()
```


**When CPU Offloading is True**
```
Initial loss: 0.859863
Initial global grad norm: 0.000000
Pre-clipping global norm: 5026.077637
Modified 4 parameter gradients
Post-clipping global norm: 5026.077637
[2025-05-21 20:30:27,023] [INFO] [logging.py:107:log_dist] [Rank 0] step=1, skipped=0, lr=[0.001], mom=[[0.9, 0.999]]
Optimizer step completed successfully

Summary:
Pre-clipping norms: [5026.07763671875]
Post-clipping norms: [5026.07763671875]
Pre-clipping global grad norm: 0.000000
Post-clipping global grad norm: 0.000000
**Final global grad norm: 1.227070**
```

**No CPU Offloading**
```
Initial loss: 0.832031
Initial global grad norm: 0.000000
Pre-clipping global norm: 4935.634191
Modified 4 parameter gradients
Post-clipping global norm: 0.761391
[2025-05-21 20:29:35,934] [INFO] [logging.py:107:log_dist] [Rank 0] step=1, skipped=0, lr=[0.001], mom=[[0.9, 0.999]]
Optimizer step completed successfully

Summary:
Pre-clipping norms: [4935.634190913677]
Post-clipping norms: [0.7613913786257984]
Pre-clipping global grad norm: 0.000000
Post-clipping global grad norm: 0.000000
**Final global grad norm: 0.000186**
```

**Problem**
Have a look at `Final global grad norm` in the outputs.
The `averaged_gradients` is not updated when CPU offloading is enabled
https://github.com/deepspeedai/DeepSpeed/blob/b048cc2b46a559b8f81c4315008cec34e650bbe0/deepspeed/runtime/zero/stage3.py#L1147

Thats why even though gradients are clipped, `_global_grad_norm` still uses the old averaged_gradients. If we change ds_config to `"offload_optimizer": {"device": "none"}`, `_global_grad_norm` will be updated.


For CPU offloading, gradients are already in `fp32_partitioned_groups_flat.grad`. Can raise a PR for this. Please let me know your thoughts @sfc-gh-truwase @tjruwase 

@therealnaveenkamal, awesome work. It looks like you really nailed the issue. A PR will be greatly appreciated. Thanks so much. 
@therealnaveenkamal Thanks for your analysis! However, when I use `offload_optimizer": {"device": "none"}`, the `_global_grad_norm` is still not updated. Is this behavior potentially related to i use bf16 mixed precision training?

Commit message: Fix: Update grad norm calculation for CPU offload

Patch to review:
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
index f37c6c271009..a13aaf7ba263 100644
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -1148,11 +1148,20 @@ def independent_gradient_partition_epilogue(self):
                     self.__param_id_to_grad_partition[param.ds_id]
                     if param.requires_grad else torch.zeros_like(param.ds_tensor) for param in sub_group
                 ]
-        # this method gets called after every backward. need to increment
-        # here because if it gets incremented in backward() the micro step
-        # id will be off by one when we do the reduce and partition at the.
-        # start of this method.
-        # TODO. make this less error prone
+        else:
+            # Update averaged_gradients to use the new fp32_partitioned_groups_flat
+            for i, sub_group in enumerate(self.fp16_groups):
+                self.averaged_gradients[i] = {}
+                for param in sub_group:
+                    if param.requires_grad:
+                        # Get the gradient from fp32_partitioned_groups_flat
+                        group_idx, dest_offset, num_elements = self.grad_position[param.ds_id]
+                        grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(
+                            0, dest_offset, num_elements)
+                        # Store the gradient in averaged_gradients
+                        self.averaged_gradients[i][param.ds_id] = grad
+                    else:
+                        self.averaged_gradients[i][param.ds_id] = torch.zeros_like(param.ds_tensor)
         self.micro_step_id += 1
 
     def overlapping_partition_gradients_reduce_epilogue(self):
@@ -1434,15 +1443,31 @@ def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param, fp32_grad_tenso
             fp32_grad_tensor.copy_(src_tensor, non_blocking=True)
             param.grad = None
 
-    def complete_grad_norm_calculation_for_cpu_offload(self, params):
+    def complete_grad_norm_calculation_for_cpu_offload(self, idx, params):
         total_norm = 0.0
         norm_type = 2.0
-        for p in params:
-            if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0):
-                param_id = self.get_param_id(p)
-                if param_id in self.norm_for_param_grads.keys():
-                    param_norm = self.norm_for_param_grads[param_id]
-                    total_norm += param_norm**2
+        
+        # First check if we have averaged gradients
+        if hasattr(self, 'averaged_gradients') and self.averaged_gradients:
+            param_grads_dict = self.averaged_gradients[idx]
+            for param in params:
+                if is_model_parallel_parameter(param) or (self.model_parallel_rank == 0):
+                    param_id = self.get_param_id(param)
+                    if param.requires_grad:
+                        param_id = self.get_param_id(param)
+                        grad_partition = param_grads_dict.get(param_id) # Get the FP32 grad partition
+
+                        if grad_partition is not None:
+                            param_norm = grad_partition.norm(norm_type)
+                            total_norm += param_norm**2
+        else:
+            # Fallback to original calculation if no averaged gradients
+            for p in params:
+                if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0):
+                    param_id = self.get_param_id(p)
+                    if param_id in self.norm_for_param_grads.keys():
+                        param_norm = self.norm_for_param_grads[param_id]
+                        total_norm += param_norm**2
 
         # Sum across all model parallel GPUs.
         total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
@@ -1882,7 +1907,7 @@ def _get_norm_groups(self):
         norm_groups = []
         for i, group in enumerate(self.fp16_groups):
             if self.offload_optimizer:
-                norm_groups.append(self.complete_grad_norm_calculation_for_cpu_offload(self.fp16_groups[i]))
+                norm_groups.append(self.complete_grad_norm_calculation_for_cpu_offload(i, self.fp16_groups[i]))
             else:
                 norm_groups.append(self.get_grad_norm_direct(self.averaged_gradients[i], self.fp16_groups[i]))
         return norm_groups


Code review comments:
diff_hunk: @@ -1148,11 +1148,20 @@ def independent_gradient_partition_epilogue(self):
                     self.__param_id_to_grad_partition[param.ds_id]
                     if param.requires_grad else torch.zeros_like(param.ds_tensor) for param in sub_group
                 ]
-        # this method gets called after every backward. need to increment
-        # here because if it gets incremented in backward() the micro step
-        # id will be off by one when we do the reduce and partition at the.
-        # start of this method.
-        # TODO. make this less error prone
+        else:
+            # Update averaged_gradients to use the new fp32_partitioned_groups_flat
+            for i, sub_group in enumerate(self.fp16_groups):
+                self.averaged_gradients[i] = {}
+                for param in sub_group:
+                    if param.requires_grad:
+                        # Get the gradient from fp32_partitioned_groups_flat
+                        group_idx, dest_offset, num_elements = self.grad_position[param.ds_id]
+                        grad = self.fp32_partitioned_groups_flat[group_idx].grad.narrow(
+                            0, dest_offset, num_elements)
path: deepspeed/runtime/zero/stage3.py
text: @user1:
This H2D will cause slowdown for backward.