Enabling ``sequence_parallel`` slows down training with fp16

[cavdard]

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,
)

[ptrendx]

@ksivaman Could you take a look?

[rahul003]

Hey @ptrendx and @ksivaman Any update on this?

[ksivaman]

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?

[ksivaman]

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

[rahul003]

Yes thanks for the help