Gradient not accumulated across nodes in deepspeed code

ahhyun: <p>Hello, I’m trying to make a deepspeed version of a code that worked without  deepspeed and see if the results can be replicated in deepspeed version. However, it seems our code is not working properly and hence wanted to ask for guidance.</p>
<p>TLDR:</p>
<ol>
<li><strong>Cannot access gradient from model_engine</strong> (only from the model itself which model_engine wraps)</li>
<li><strong>Gradient does not seem to be aggregated across gpus after <code>.backward</code></strong> (i.e. when we run with 4 gpus with batch size of 32 each, the batch size of the input is 32 while the batch size of the gradient is also 32)</li>
<li>The validation (eval mode) loss for the same data is the same for all ranks except rank 0.</li>
</ol>
<p>We would greatly appreciate if anyone could give us even a slight hint into debugging this.</p>
<p>Below is the a detailed explanation.</p>
<h3><a name="p-472621-background-1" class="anchor" href="#p-472621-background-1"></a>Background</h3>
<p>Deepspeed version and no deepspeed version code exists separately.</p>
<ul>
<li>deepspeed version code: <code>pretrain_main_deepspeed.py</code>, <code>pretrain_trainer_deepspeed.py</code></li>
<li>no deepspeed version code : <code>pretrain_main.py</code>, <code>pretrain_trainer.py</code></li>
</ul>
<p>Experimental conditions are</p>
<ul>
<li>IDENTICAL :  global batch size, lr, wd, similar optimizer(AdamW , FusedAdam)</li>
<li>DIFFERENT : number of gpus<br>
=&gt; We expect them to give same training curves.</li>
</ul>
<p>I found out that 1node 4GPU version(deepspeed) seem to work fine based on validation loss, following the loss curve of 1node 4gpu(torch’s Data Parallel). (Orange : deepspeed, Blue: no deepspeed)</p>
<p><img width="690" height="411" alt="Image" src="https://github.com/user-attachments/assets/72d9db5e-4f63-496b-b505-8c662327219a"></p>
<p>However, in 16 node condition, loss doesn’t converge to expected level. I’ve tried variants of below.</p>
<ul>
<li>(a) same effective batch size : batch_size 2 (effective batch size = 2 * 16Nodes * 4GPUs = 128)</li>
<li>(b) batch_size(same batch_size per GPU : 128 / batch_size per GPU : 32) X lr (optimal lr in 1node condition * <strong>N</strong> / *<strong>sqrt(N)</strong>)</li>
</ul>
<h3><a name="p-472621-implementation-details-2" class="anchor" href="#p-472621-implementation-details-2"></a>Implementation details</h3>
<p>I’ve attached the script I’ve used.<br>
Note) <code>train_micro_batch_size_per_gpu</code> is not used in the actual code, argument <code>batch_size</code> is the actual batch size fed into dataloader.</p>
<pre><code class="lang-auto">    data_loader = DataLoader(
        train_dataset, batch_size=params.batch_size, # This is per-GPU batch size
        sampler=train_sampler,
        num_workers=16, shuffle=(train_sampler is None), # Shuffle is True if not using DDP sampler
        persistent_workers=True,
        prefetch_factor=3, pin_memory=True, drop_last=True, collate_fn=collate_fn_for_data_info
    )
</code></pre>
<p><a href="https://drive.google.com/file/d/1Z32fQ4EqHYaOK69evavCuWmoZeX8uAQ2/view?usp=share_link" rel="noopener nofollow ugc">Click to download code used</a></p>
<h3><a name="p-472621-problem-questions-3" class="anchor" href="#p-472621-problem-questions-3"></a>Problem &amp; Questions</h3>
<p>So I’m guessing there’s fundamental problem with my deepspeed code.<br>
I’ve printed out below to see if gradient is accumulated across ranks.<br>
There were three problems</p>
<ol>
<li>Cannot inspect gradient from <code>model_engine</code></li>
<li>Gradient shape not as expected</li>
<li>Validation loss with fake data gave different loss only for rank 0.</li>
</ol>
<h4><a name="p-472621-questions-4" class="anchor" href="#p-472621-questions-4"></a>Questions</h4>
<ol>
<li>
<p>How to inspect gradient across ranks?<br>
I’ve tried two versions as below. <code>model</code> printed out gradients while <code>model_engine</code> didn’t. (2node 4GPU)</p>
<pre><code>def debug_gradients(self, batch_idx, log_every=10):
     """Quick gradient debugging"""
     if batch_idx % log_every != 0:
         return

 rank = deepspeed.comm.get_rank() if self.params.deepspeed else 8888
 print(f"\n[Rank {rank}] Gradient Check - Batch {batch_idx}")
 
 print("✅MODEL VERSION")
 for name, param in self.model.named_parameters():
     grad = deepspeed.utils.safe_get_full_grad(param)
     if grad is not None:
         print("Gradient Shape: ", name, grad.shape)
     else:
         print("Gradient is None for parameter:", name)
         
 print("✅MODEL ENGINE VERSION")       
 for name, param in self.model_engine.named_parameters():
     if param.grad is not None:
         grad_norm = param.grad.norm().item()
         print(f"  {name}: shape={param.grad.shape}, norm={grad_norm:.4f}")
     else:
         print(f"  {name}: No gradient")`
</code></pre>
</li>
</ol>
<pre><code class="lang-auto">
x.shape:  torch.Size([32, 19, 30, 500])

✅MODEL VERSION
Gradient is None for parameter: mask_encoding
Gradient Shape:  embedding.0.proj_in.0.weight torch.Size([32, 1, 1, 63])
Gradient Shape:  embedding.0.proj_in.0.bias torch.Size([32])
Gradient Shape:  embedding.0.proj_in.1.weight torch.Size([32])
Gradient Shape:  embedding.0.proj_in.1.bias torch.Size([32])
Gradient Shape:  embedding.0.proj_in.3.weight torch.Size([32, 32, 1, 3])

✅MODEL ENGINE VERSION
  module.mask_encoding: No gradient
  module.embedding.0.proj_in.0.weight: No gradient
  module.embedding.0.proj_in.0.bias: No gradient
  module.embedding.0.proj_in.1.weight: No gradient
  module.embedding.0.proj_in.1.bias: No gradient
</code></pre>
<ol start="2">
<li>
<p>Also, input shape, gradient shape weren’t as I expected.<br>
(1node 4GPU)<br>
Since batch_size is 32, I think gradient shape should be (128, -,-,-) since gradient should be shared across model_engine. How should I target this?</p>
</li>
<li>
<p>I’ve checked validation loss with fake data (same shape with input x, <code>torch.ones</code> and <code>torch.zeros</code>) and found validation loss differ only in rank 0. (2node 4GPU)<br>
</p>
</li>
</ol>
<pre><code class="lang-auto">def validate(self, epoch, normalize_factor=100.0):
        self.model_engine.eval()
        valid_losses_rank = [] # Losses on this rank

        # tqdm only on rank 0
        iterable_valid_loader = self.valid_data_loader
        if self.is_rank_0:
            iterable_valid_loader = tqdm(self.valid_data_loader, desc=f"Validation Epoch {epoch}", mininterval=10)

        with torch.no_grad():
            for batch_idx, (x, data_info_list) in enumerate(iterable_valid_loader):
                x = x.to(self.device, dtype=torch.float32) / 100.0
                print("validation x.shape: ", x.shape)
                
                ##!
                fake_data = torch.ones_like(x, device=self.device, dtype=torch.float32)/2 # Create (same) fake data across ranks
                fake_loss = self.SSL.compute_loss(fake_data, data_info_list=data_info_list)
                print(f"Fake loss for validation for rank {deepspeed.comm.get_rank()}: {fake_loss.item()}") # Print fake loss for debugging
                ##!

</code></pre>
<pre><code class="lang-auto">Fake loss for validation for rank 0: 0.8215165138244629
Fake loss for validation for rank 4: 0.8218275308609009
Fake loss for validation for rank 5: 0.8218275308609009
Fake loss for validation for rank 7: 0.8218275308609009
Fake loss for validation for rank 6: 0.8218275308609009
Fake loss for validation for rank 2: 0.8218275308609009
Fake loss for validation for rank 1: 0.8218275308609009
Fake loss for validation for rank 3: 0.8218275308609009
</code></pre>

yf225: <p>I believe you are trying to debug deepspeed code. Maybe this should be posted in deepspeed github issues / forum?</p>

ahhyun: <p>I will!<br>
Thank you for your reply.</p>