Skip to content

Enabling sequence_parallel slows down training with fp16 #182

Closed
Closed
Enabling sequence_parallel slows down training with fp16#182
Assignees
ksivaman

Description

@cavdard
cavdard
opened on Apr 27, 2023

I am testing GPT2 model training using TransformerLayer.

Training slows down significantly when sequence_parallel=True, achieves 1/5th of throughput of training without sequence_parallel.
I also observe that sequence_parallel=True results in OOM for some batch sizes where sequence_parallel=False can run successfully.

Do you have any recommendation to achieve better throughput with sequence_parallel and fp16?

Model is ~4.3B with 12 layers, tp_size=4, fp16, seq_len=2048, training with 8 A100 GPUs.

transformer_engine.pytorch.TransformerLayer(
5120,
20480,
40,
layer_number=(l+1),
self_attn_mask_type="causal",
tp_group=tp_group(),
tp_size=tp_size,
params_dtype=torch.float16,
output_layernorm=True,
layer_type="encoder",
set_parallel_mode=True,
fuse_qkv_params=True,
sequence_parallel=True,
qkv_weight_interleaved=False,
attention_softmax_in_fp32=False,
)

Activity

  1. ptrendx
    assigned
    ksivaman
    on Apr 28, 2023

  2. ptrendx commented on Apr 28, 2023

    @ptrendx
    ptrendx
    on Apr 28, 2023
    Member

    @ksivaman Could you take a look?

  3. rahul003 commented on May 9, 2023

    @rahul003
    rahul003
    on May 9, 2023
    Contributor

    Hey @ptrendx and @ksivaman Any update on this?

  4. ksivaman commented on May 9, 2023

    @ksivaman
    ksivaman
    on May 9, 2023
    Member

    Which toolkit (e.g. NeMo) are you using to enable sequence parallelism? Does that toolkit support sequence parallel? Simply passing sequence_parallel=True in TE will not do the trick. The input during the forward pass must also be split in the sequence dimension among the tensor parallel workers before passing it to TE's TransformerLayer.

    If this is the case, an unsharded input will be gathered across the TP group during fprop, making the input larger and explaining the extra memory consumption as well as slowdown observed. I think this is what is happening here, does that make sense?

  5. ksivaman commented on May 9, 2023

    @ksivaman
    ksivaman
    on May 9, 2023
    Member

    Could you confirm if this is the issue? Either way, I will make a change soon that catches this behavior and reports an error.

  6. rahul003 commented on Jun 5, 2023

    @rahul003
    rahul003
    on Jun 5, 2023
    Contributor

    Yes thanks for the help

  7. NVIDIA
    locked and limited conversation to collaborators on Jun 16, 2023
  8. ksivaman
    converted this issue into a discussion #283 on Jun 16, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions