Question about GPU memory usage when using pipeline parallelism training under larger micro batch count

wangmingyuan: <p>I am using torchtian with FSDP2 + PP(1F1B) to train llama3-8b, however, I found that as the <code>micro batch count</code> increasing, the GPU memory usage will increase rapidly (from 42.27GB to 64.57GB) on the last pp stage. That’s a bit strange, AFAK, the GPU memory should not increase too much though use larger <code>micro batch count</code> (After all, we need to use larger micro batch count to decrease the bubble rate).</p>
<p>Here is my experiment settings and results:</p>
<p>using 4 GPUs and DP2-PP2 to train llama3-8b (pruned, 16 layers)<br>
pytorch version: 2.9.0.dev20250725+cu126<br>
torchtitan version: v0.1.0</p>
<p>command:</p>
<pre data-code-wrap="shell"><code class="lang-shell">LOG_RANK=2 NGPU=4 CONFIG_FILE=./torchtitan/models/llama3/train_configs/llama3_8b.toml ./run_train.sh --metrics.log_freq 1  --training.seq_len 4096 --training.steps 1000 --parallelism.data_parallel_shard_degree 2 --activation_checkpoint.mode full --parallelism.pipeline_parallel_degree 2  --training.local_batch_size 16
</code></pre>
<p>memory: torch max reserved memory<br>
rank0 is pp stage0, rank2 is pp stage1 (last stage)</p>
<div class="md-table">
<table>
<thead>
<tr>
<th>micro batch count</th>
<th>memory (GB) on rank0</th>
<th>memory (GB) on rank2</th>
</tr>
</thead>
<tbody>
<tr>
<td>2</td>
<td>35.75</td>
<td>40.30</td>
</tr>
<tr>
<td>4</td>
<td>35.80</td>
<td>42.27</td>
</tr>
<tr>
<td>8</td>
<td>35.94</td>
<td>48.27</td>
</tr>
<tr>
<td>16</td>
<td>36.20</td>
<td>65.47</td>
</tr>
</tbody>
</table>
</div><p>Also, I dumped the memory snapshot of last PP stage and found that the output (shape: 4096 * 128256) of LM head in each micro batch keeps alive during the step period<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/d/3d911822a7292ad20fa5eac134ba2e1ba46ded40.png" data-download-href="https://discuss.pytorch.org/uploads/default/3d911822a7292ad20fa5eac134ba2e1ba46ded40" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/d/3d911822a7292ad20fa5eac134ba2e1ba46ded40_2_689x285.png" alt="image" data-base62-sha1="8MDZiMedLirEDjHXT0puo0deAGA" width="689" height="285" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/d/3d911822a7292ad20fa5eac134ba2e1ba46ded40_2_689x285.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/d/3d911822a7292ad20fa5eac134ba2e1ba46ded40_2_1033x427.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/d/3d911822a7292ad20fa5eac134ba2e1ba46ded40_2_1378x570.png 2x" data-dominant-color="D7CCC4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">2536×1049 303 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I checked the implementation of pipeline parallelism and found that PyTorch will caching chunk outputs after <code>forward_maybe_with_no_sync</code> in <code>forward_one_chunk</code> function, this may explains the GPU memory behavior of the above, but <strong>my question is: Do we really need to cache the chunk output during training</strong>, perhaps we only need to cache output during evaluation or inference ? Or let the users choose whether enable caching outputs or not.</p>
<p>After I tried to disable caching outputs with some modifications like below, the GPU memory usage looks normal:</p>
<pre data-code-wrap="python"><code class="lang-python">class _PipelineStageBase(ABC):
    def __init__(self, ...):
        ...
        # use `requires_merge_outputs` as a switch
        self.requires_merge_outputs: bool = True
    
    def forward_one_chunk(self, ...):
        ...
        output = self.forward_maybe_with_nosync(*composite_args, **composite_kwargs)
        if self.is_last and self.requires_merge_outputs:
            self.output_chunks.append(output)
        ...

class PipelineScheduleSingle(_PipelineSchedule):
    def step(self, *args, target=None, losses: Optional[list] = None, **kwargs):
        ...
        if self._stage.is_last and self._stage.requires_merge_outputs:
            return self._merge_outputs(self._stage.output_chunks)
        else:
            return None
    
    # expose to the user
    def set_requires_merge_outputs(self, requires_merge_outputs: bool):
        self._stage.requires_merge_outputs = requires_merge_outputs
</code></pre>
<p>Now the GPU memory usage more stable</p>
<div class="md-table">
<table>
<thead>
<tr>
<th>micro batch count</th>
<th>memory (GB) on rank0</th>
<th>memory (GB) on rank2</th>
</tr>
</thead>
<tbody>
<tr>
<td>2</td>
<td>35.75</td>
<td>40.30</td>
</tr>
<tr>
<td>4</td>
<td>35.80</td>
<td>40.38</td>
</tr>
<tr>
<td>8</td>
<td>35.94</td>
<td>40.41</td>
</tr>
<tr>
<td>16</td>
<td>36.20</td>
<td>41.64</td>
</tr>
</tbody>
</table>
</div>

H-Huang: <p>If we don’t save the outputs for the last PP stage, then the memory bottleneck will probably be in them middle stages so the memory required may not differ that much. Have you tested to see what is the max memory allocated for the other ranks?</p>
<p>You still bring up a good point though and that feature seems like something we should support. We have a new <a href="https://github.com/pytorch/pytorch/pull/157795" rel="noopener nofollow ugc"><code>eval()</code> API</a>, so users can use that instead if they need the outputs saved. Can you file an issue on github? <a href="https://github.com/pytorch/pytorch/issues" class="inline-onebox" rel="noopener nofollow ugc">GitHub · Where software is built</a>, we can better track it there</p>
<p>EDIT: nvm I see you already filed an issue, thanks!</p>

wangmingyuan: <p><a class="mention" href="/u/h-huang">@H-Huang</a> I supplemented the max reserved GPU memory on other ranks(or stages) in the previous table, it seems that whether cache the forward outputs on last PP stage or not has no effect on the memory footprint of other ranks.</p>
<p>By the way, you indicated that “If we don’t save the outputs for the last PP stage, then the memory bottleneck will probably be in them middle stages”. I find this point a bit confusing,  could you help me understand the reasoning behind this in a bit more detail when you have time.</p>
<p>Thanks in advance !</p>

H-Huang: <p>I don’t mean to say that the caching will affect the memory of the middle stages, I just wonder how the memory profile of the entire training job now looks.</p>
<p>The max memory usage should be calculated as max(all_ranks) since if any one of the ranks OOMs then the training will fail. So I was just wondering how the memory usage of the middle ranks look like. If last stage was the previous bottleneck where is the new bottleneck, was just curious <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=14" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

wangmingyuan: <p>Got it, thanks for clarifying!</p>