Official example for multi-GPU training didn't work on 4090

[QIU-Shuo]

Hi, I am working on the multi-GPU training example here with nvidia's pytorch docker 23.07 (without changing anything), and a quad GPU machine. However I got the error below. Is there a way to resolve it?

Traceback (most recent call last):
  File "/home/qiushuo/jipiti/te_bench.py", line 56, in <module>
    utils.speedometer(
  File "/home/qiushuo/jipiti/quickstart_utils.py", line 33, in speedometer
    output = module(input, **forward_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/parallel/distributed.py", line 1513, in forward
    inputs, kwargs = self._pre_forward(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/parallel/distributed.py", line 1401, in _pre_forward
    if torch.is_grad_enabled() and self.reducer._rebuild_buckets():
RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by passing the keyword argument `find_unused_parameters=True` to `torch.nn.parallel.DistributedDataParallel`, and by 
making sure all `forward` function outputs participate in calculating loss. 
If you already have done the above, then the distributed data parallel module wasn't able to locate the output tensors in the return value of your module's `forward` function. Please include the loss function and the structure of the return value of `forward` of your module when reporting this issue (e.g. list, dict, iterable).
Parameter indices which did not receive grad for rank 0: 2 3 4 5 6 7
 In addition, you can set the environment variable TORCH_DISTRIBUTED_DEBUG to either INFO or DETAIL to print out information about which particular parameters did not receive gradient on this rank as part of this error

my launch commnad:

torchrun --nproc_per_node 1 --nnodes 1 --rdzv_id=100 --rdzv_backend=c10d --rdzv_endpoint=localhost:29312 te_bench.py

My script (only changed a few size parameter to acoomodate flash attention requirement):

import torch
import transformer_engine.pytorch as te
from transformer_engine.common.recipe import Format, DelayedScaling
import quickstart_utils as utils
import os

# Layer configuration
hidden_size = 2048
sequence_length = 2048
batch_size = 4
ffn_hidden_size = 11008
num_attention_heads = 32
dtype = torch.float16

# Synthetic data
x = torch.rand(sequence_length, batch_size, hidden_size).cuda().to(dtype=dtype)
dy = torch.rand(sequence_length, batch_size, hidden_size).cuda().to(dtype=dtype)

fp8_format = Format.HYBRID
fp8_recipe = DelayedScaling(
    fp8_format=fp8_format,
    amax_history_len=16,
    amax_compute_algo="max",
)
world_group = torch.distributed.init_process_group(
    "nccl",
    init_method="file:///tmp/rdzv",
    world_size=1,
    rank=0,
)
data_parallel_group = torch.distributed.new_group(ranks=[0], backend="nccl")
tensor_parallel_group = torch.distributed.new_group(ranks=[0], backend="nccl")


# Construct layer
parallel_transformer = te.TransformerLayer(
    hidden_size,
    ffn_hidden_size,
    num_attention_heads,
    set_parallel_mode=True,
    tp_group=tensor_parallel_group,
    sequence_parallel=True,
)
parallel_transformer.to(dtype=dtype).cuda()
parallel_transformer = torch.nn.parallel.DistributedDataParallel(
    parallel_transformer,
    process_group=data_parallel_group,
)

# Training step
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe, fp8_group=data_parallel_group):
    y = parallel_transformer(x, attention_mask=None)
y.backward(dy)

# Measure step time
utils.speedometer(
    parallel_transformer,
    x,
    dy,
    forward_kwargs = { "attention_mask": None },
    fp8_autocast_kwargs = {
        "enabled": True,
        "fp8_recipe": fp8_recipe,
        "fp8_group": data_parallel_group,
    },
)

[Xiaoyiyong555]

hi i also meet this problem! have you solve it?