Lightning-AI / Lightning-AI/litgpt

Bug: Incorrect gradient accumulation steps calculation in multi-node training due to missing world size information

Open
#1,927 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
Dominant language
Python
Stars
13.7k
Forks
1.5k
Avg merge
15h 37m
Merged PRs (30d)
1

Description

### Bug description

There's a bug in the gradient accumulation calculation that affects multi-node training scenarios. The current implementation uses local device count instead of world size (total devices across all nodes), leading to incorrect gradient accumulation steps and training behavior.

**Current Behavior**
The `gradient_accumulation_iters` function calculates steps based on local batch size:
```python
def gradient_accumulation_iters(self, devices: int) -> int:
"""Number of iterations between gradient synchronizations"""
gradient_accumulation_iters = self.batch_size(devices) // self.micro_batch_size
return gradient_accumulation_iters

def batch_size(self, devices: int) -> int:
"""Number of samples between optimizer steps per data-parallel rank"""
batch_size = self.global_batch_size // devices # devices is local count only
return batch_size
```

The `devices` parameter comes from `torch.cuda.device_count()`, which only returns local GPU count (e.g., 8) rather than total GPUs across all nodes (e.g., 8 * num_nodes = 128).

**Impact**
This causes:
1. Incorrect gradient accumulation frequency (off by a factor of num_nodes)
2. Mismatch between steps and iterations in training logs
3. Example from logs:
```
Epoch 1 | iter 109472 step 13684 | loss train: 3.099, val: 3.075
```
Note the 8x difference between iterations and steps (109472/13684 ≈ 8)

**Steps to Reproduce**
1. Configure multi-node training (e.g., 16 nodes, 8 GPUs each)
2. Set global batch size and micro batch size
3. Observe gradient accumulation steps and training logs

**Expected Behavior**
The calculation should use total world size (all devices across all nodes) instead of local device count:
```python
def batch_size(self, devices: int) -> int:
"""Number of samples between optimizer steps per data-parallel rank"""
batch_size = self.global_batch_size // fabric.world_size # Use world_size instead of devices
return batch_size
```

**Proposed Solution**
Use `fabric.world_size` instead of local `devices` count to properly account for all processes across nodes.

**Additional Context**
- This issue likely wasn't caught earlier because the tutorials primarily use single-node setups (e.g., https://github.com/Lightning-AI/litgpt/blob/a5021be4bb48e27779586b56b062a1749ecb232f/config_hub/pretrain/tinyllama.yaml#L109)
- The bug becomes apparent only in multi-node training scenarios
- This affects training convergence and effective learning rate in multi-node setups

Let me know if you'd like me to modify any part of this issue description before you post it.

### What operating system are you using?

Linux

### LitGPT Version

```

```

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the gradient_accumulation_iters and batch_size entry points described in the issue, then trace how torch.cuda.device_count() supplies the devices value and how Fabric exposes world size. Reproduce the 16-node, 8-GPU scenario using the linked tinyllama.yaml training configuration and verify that accumulation frequency and iteration-to-step ratios match the total world size.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.