Help with Gradient and Optimizer Step for Tensor Parallel Implementations

username: Qikai_Lu
cooked: <p>I have recently implemented my own transformer-TP, and have noticed that compared to single GPU training, 2-GPU training yielded significantly worse perplexity score. I suspect that this is caused by incorrect gradient synchronization, located somewhere between backward.loss() and optimizer.step(). So I am wondering if there are any pytorch TP tutorial that shows proper loss computation and weight computation? Alternatively, if there are any tips on how to perform loss compute and parameter updates on TP-model, that would be helpful as well.</p>
<p>For reference, this is my model loss compute and backward, and optimizer step(). It is preferable to not use loss_parallel (which only support cross-entropy), as I intend to add some regularization terms and secondary losses thereafter.</p>
<pre><code class="lang-auto">        # Part of the transformer TP plan
        for trsfm_block in self.encoder.layers:
            parallelize_plan = {
                "norm_attn": SequenceParallel(),
                "attn_block": PrepareModuleInput(input_kwarg_layouts={"x": Shard(1),
                                                                      "causal_mask": None},
                                                 desired_input_kwarg_layouts={"x": Replicate(),
                                                                              "causal_mask": None}),
                "attn_block.w_q": ColwiseParallel(),
                "attn_block.w_k": ColwiseParallel(),
                "attn_block.w_v": ColwiseParallel(),
                "attn_block.w_o": RowwiseParallel(output_layouts=Shard(1)),
                "norm_ff": SequenceParallel(),
                "ff_block": PrepareModuleInput(input_kwarg_layouts={"x": Shard(1)},
                                               desired_input_kwarg_layouts={"x": Replicate()}),
                "ff_block.linear_1": ColwiseParallel(),
                "ff_block.linear_2": RowwiseParallel(output_layouts=Shard(1)),
            }
            parallelize_module(
                module=trsfm_block,
                device_mesh=self.proc_group_info["mesh_common"],
                parallelize_plan=parallelize_plan
            )
            trsfm_block.attn_block.n_heads = trsfm_block.attn_block.n_heads // self.proc_group_info["mesh_common"].size()
        ......
        # Compute loss after forward pass.
        if label is not None:
            # Note that input is the replicated transformer output tensor.
            loss = self.xent_loss(input=x.view(-1, x.size(-1)), target=label.view(-1))
            loss.backward()
            # Should gradient sync be done here?
            optimizer.step()
</code></pre>
<p>For the official PyTorch TP (Tensor Parallel) examples:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/tensor_parallel_example.py">
  <header class="source">

      <a href="https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/tensor_parallel_example.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/tensor_parallel_example.py" target="_blank" rel="noopener nofollow ugc">pytorch/examples/blob/main/distributed/tensor_parallelism/tensor_parallel_example.py</a></h4>


      <pre><code class="lang-py">import os
import sys
import torch
import torch.nn as nn

from torch.distributed.tensor.parallel import (
    parallelize_module,
    ColwiseParallel,
    RowwiseParallel,
)

from log_utils import rank_log, get_logger, verify_min_gpu_count

# ---- GPU check ------------
_min_gpu_count = 2

if not verify_min_gpu_count(min_gpus=_min_gpu_count):
    print(f"Unable to locate sufficient {_min_gpu_count} gpus to run this example. Exiting.")
    sys.exit()
# ---------------------------
</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/tensor_parallel_example.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/fsdp_tp_example.py">
  <header class="source">

      <a href="https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/fsdp_tp_example.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/fsdp_tp_example.py" target="_blank" rel="noopener nofollow ugc">pytorch/examples/blob/main/distributed/tensor_parallelism/fsdp_tp_example.py</a></h4>


      <pre><code class="lang-py">import sys
import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F

from log_utils import rank_log, get_logger, verify_min_gpu_count

# ---- GPU check ------------
_min_gpu_count = 4

if not verify_min_gpu_count(min_gpus=_min_gpu_count):
    print(f"Unable to locate sufficient {_min_gpu_count} gpus to run this example. Exiting.")
    sys.exit()
# ---------------------------

from llama2_model import Transformer, ModelArgs

from torch.distributed.device_mesh import init_device_mesh
</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/fsdp_tp_example.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
I noticed that the parameter update is incomplete (there is not label), as they are meant as tutorials.</p>

username: fduwjj
cooked: <p>Here is the tutorial for TP in PyTorch: <a href="https://pytorch.org/tutorials/intermediate/TP_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">Large Scale Transformer model training with Tensor Parallel (TP) — PyTorch Tutorials 2.5.0+cu124 documentation</a>. We do the parameter update inside TP API.</p>