[BUG]: NaN loss while using Gemini + tensor_parallel

[airlsyn]

🐛 Describe the bug

Using llama2 benchmark.py with the change:

if args.plugin == "gemini":
        plugin = GeminiPlugin(
            precision="bf16",
            shard_param_frac=args.shard_param_frac,
            offload_optim_frac=args.offload_optim_frac,
            offload_param_frac=args.offload_param_frac,
            tp_size=args.tp,
        )
....

1. plugin = '3d' + tp=2
   loss is normal

torchrun --nnodes=1 --nproc_per_node=4 --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=localhost:12345 \
  ColossalAI/examples/language/llama2/benchmark.py \
  -c 7b -g -b 1 --max_length 4096 -p 3d --tp 2 --pp 1 --zero 2 -s 100

2. plugin = 'gemini'
   loss is normal

torchrun --nnodes=1 --nproc_per_node=4 --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=localhost:12345 \
  ColossalAI/examples/language/llama2/benchmark.py \
  -c 7b -g -b 1 --max_length 4096 -p gemini --tp 1 --pp 1 --zero 2 -s 100

3. plugin = 'gemini' + tp=2
   loss is normal at the beginning, and changes to NaN later

torchrun --nnodes=1 --nproc_per_node=4 --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=localhost:12345 \
  ColossalAI/examples/language/llama2/benchmark.py \
  -c 7b -g -b 1 --max_length 4096 -p gemini --tp 2 --pp 1 --zero 2 -s 100

Why loss goes to NaN when using 'gemini' + tensor parallel? And how to fix it?

Thanks a lot.

Environment

Pytorch: 2.0.1
ColossalAI: main branch (CUDA_EXT=1 pip install .)

[airlsyn]

Hello, @flybird11111 , do you have a moment to take a look? Thanks so much.

[flybird11111]

Hello, @flybird11111 , do you have a moment to take a look? Thanks so much.

Sure, I will take a look at it.

[airlsyn]

Hello, @flybird11111 , do you have a moment to take a look? Thanks so much.

Sure, I will take a look at it.

Thanks very much. May I ask "have you figure it out"?

[flybird11111]

Hello, @flybird11111 , do you have a moment to take a look? Thanks so much.

Sure, I will take a look at it.

Thanks very much. May I ask "have you figure it out"?

The issue is still being debugged.

[flybird11111]

Hi, this issue has been resolved.

[airlsyn]

Hi, this issue has been resolved.

Thanks a lot @flybird11111, it works fine for randomly generated data

However, when using real data such as ThePile. I observe numerous zero values in input_embeds a few steps later.

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

. It appears like the following: Zero values consistently appear at either the beginning or the end.

If I comment VocabParallelEmbedding1D replacement, thing looks like well

            self.append_or_create_submodule_replacement(
                description=SubModuleReplacementDescription(
                    suffix="embed_tokens",
                    target_module=VocabParallelEmbedding1D,
                ),
                policy=policy,
                target_key=LlamaModel,
            )

[flybird11111]

Thanks, I'll look into it.

[airlsyn]

The loss becomes irregular when VocabParallelEmbedding1D is not utilized in gemini + tp=8 and hybrid + tp=8 + zero2, in comparison to gemini without tp

[airlsyn]

@flybird11111 May I ask have you figure it out? Thanks so much.

[airlsyn]

when using Gemini + tp_size=2 (4gpus), I found VocabParallelEmbedding1D.forward gave me a weird output after reduce_forward

ColossalAI/colossalai/shardformer/layer/embedding.py

Line 315 in af95267

output = reduce_forward(embedding_output, self.process_group)

numerous Nan value of embedding output in step 11

[pretrain.py:main] - step: 10, loss: 11.12129
tensor([[[ 0.0101,  0.0159,  0.0010,  ..., -0.0123,  0.0070, -0.0052],
         [-0.0068, -0.0219, -0.0227,  ...,  0.0179,  0.0201,  0.0330],
         [ 0.0131, -0.0005, -0.0194,  ..., -0.0024,  0.0483,  0.0092],
         ...,
         [ 0.0011,  0.0023,  0.0153,  ...,  0.0184, -0.0029, -0.0283],
         [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
         [-0.0481,  0.0070,  0.0270,  ...,  0.0008, -0.0233, -0.0151]]],
       device='cuda:3', grad_fn=<_ReduceForwardBackward>)
tensor([[[ 0.0101,  0.0159,  0.0010,  ..., -0.0123,  0.0070, -0.0052],
         [-0.0068, -0.0219, -0.0227,  ...,  0.0179,  0.0201,  0.0330],
         [ 0.0131, -0.0005, -0.0194,  ..., -0.0024,  0.0483,  0.0092],
         ...,
         [ 0.0011,  0.0023,  0.0153,  ...,  0.0184, -0.0029, -0.0283],
         [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
         [-0.0481,  0.0070,  0.0270,  ...,  0.0008, -0.0233, -0.0151]]],
       device='cuda:2', grad_fn=<_ReduceForwardBackward>)
tensor([[[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [ 0.0054, -0.0106, -0.0187,  ...,  0.0071,  0.0410,  0.0117],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [ 0.0070, -0.0264,  0.0535,  ..., -0.0254, -0.0137,  0.0337],
         [ 0.0327,  0.0383, -0.0133,  ...,  0.0101,  0.0110,  0.0547],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]]],
       device='cuda:1', grad_fn=<_ReduceForwardBackward>)
tensor([[[    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         [ 0.0054, -0.0106, -0.0187,  ...,  0.0071,  0.0410,  0.0117],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan],
         ...,
         [ 0.0070, -0.0264,  0.0535,  ..., -0.0254, -0.0137,  0.0337],
         [ 0.0327,  0.0383, -0.0133,  ...,  0.0101,  0.0110,  0.0547],
         [    nan,     nan,     nan,  ...,     nan,     nan,     nan]]],
       device='cuda:0', grad_fn=<_ReduceForwardBackward>)
[pretrain.py:main] - step: 11, loss: nan