Adding the new feature of FPDT (#441)

[saforem2]

• pass batch_dim_idx to deepspeed sequence parallel distributed attention for supporting batch size larger than 1

• add FPDT support; add Ulysses rotary position embedding support

• add FPDT support; add Ulysses rotary position embedding support

• add FPDT support; add Ulysses rotary position embedding support

• add FPDT support; add Ulysses rotary position embedding support

• remove unnecessary files

• set the warmup length to be FPDT chunk size if enabled

---

Summary by Sourcery

Add support for FPDT to improve model performance with large batch sizes and introduce Ulysses rotary position embedding for better positional encoding. Include a new script for pretraining GPT-3 models with FPDT, and optimize warmup length to match FPDT chunk size.

New Features:

• Introduce FPDT (Fast Parallel Distributed Transformer) support to enhance model performance with large batch sizes.
• Add Ulysses rotary position embedding support for improved positional encoding.

Enhancements:

• Set the warmup length to match the FPDT chunk size when enabled, optimizing initialization.

Build:

• Add a new script for pretraining GPT-3 models with FPDT support, including configuration for various model sizes and training parameters.

[sourcery-ai]

Reviewer's Guide by Sourcery

This PR implements FPDT (Fast Path Decoupled Training) support and Ulysses rotary position embedding in the transformer model. The main changes include adding FPDT functionality for sequence parallel processing, modifying the attention mechanism to support FPDT, and updating the rotary position embedding implementation.

Sequence diagram for FPDT in Transformer forward pass

sequenceDiagram
    participant Transformer
    participant FPDT_FFN
    participant DenseHTo4H
    participant ActivationFunc
    participant Dense4HToH
    Transformer->>FPDT_FFN: apply(hidden_states, dense_h_to_4h.weight, dense_h_to_4h.bias, dense_4h_to_h.weight, dense_4h_to_h.bias, add_bias, fpdt_FFN_chunk_size)
    alt FPDT not enabled
        Transformer->>DenseHTo4H: dense_h_to_4h(hidden_states)
        DenseHTo4H-->>Transformer: intermediate_parallel, bias_parallel
        Transformer->>ActivationFunc: activation_func(intermediate_parallel)
        ActivationFunc-->>Transformer: intermediate_parallel
        Transformer->>Dense4HToH: dense_4h_to_h(intermediate_parallel)
        Dense4HToH-->>Transformer: output, output_bias
    end
    Transformer-->>FPDT_FFN: output, output_bias

Loading

Sequence diagram for Rotary Position Embedding

sequenceDiagram
    participant ParallelTransformerLayer
    participant RotaryEmbedding
    participant ApplyRotaryPosEmb
    ParallelTransformerLayer->>RotaryEmbedding: rotary_pos_emb(seq_length)
    RotaryEmbedding-->>ParallelTransformerLayer: rotary_pos_emb_cos, rotary_pos_emb_sin
    ParallelTransformerLayer->>ApplyRotaryPosEmb: apply_rotary_pos_emb(query_layer, rotary_pos_emb_cos)
    ParallelTransformerLayer->>ApplyRotaryPosEmb: apply_rotary_pos_emb(key_layer, rotary_pos_emb_sin)

Loading

Class diagram for FPDT and Rotary Position Embedding

classDiagram
    class Transformer {
        +bool ds_sequence_parallel_fpdt
        +int fpdt_FFN_chunk_size
        +forward(hidden_states)
    }
    class ParallelTransformerLayer {
        +bool ds_sequence_parallel_fpdt
        +bool ds_sequence_parallel_fpdt_offloading
        +self_attention
        +forward(hidden_states, attention_mask, inference_params, rotary_pos_emb)
    }
    class RotaryEmbedding {
        +Tensor inv_freq
        +float theta
        +forward(max_seq_len, offset)
    }
    class FPDT_Attention {
        +Parameter qkv_linear_weight
        +Parameter qkv_linear_bias
        +Parameter qkv_dense_weight
        +Parameter qkv_dense_bias
        +FPDT_Attention(config, qkv_linear_weight, qkv_linear_bias, qkv_dense_weight, qkv_dense_bias, sequence_process_group, gather_idx, return_bias, chunk_size, enable_offloading)
    }
    Transformer --> FPDT_Attention
    ParallelTransformerLayer --> FPDT_Attention
    ParallelTransformerLayer --> RotaryEmbedding

Loading

File-Level Changes

Change	Details	Files
Added FPDT (Fast Path Decoupled Training) support for sequence parallel processing	Added FPDT configuration options and parameters Implemented FPDT chunk size handling for sequence parallelism Added FPDT FFN and Attention implementations Modified warmup function to use FPDT chunk size when enabled	megatron/model/transformer.py megatron/arguments.py megatron/initialize.py
Modified attention mechanism to support FPDT and batch processing	Updated query/key/value tensor handling for FPDT Added support for batch dimension index in distributed attention Modified attention mask and bias handling for FPDT compatibility	megatron/model/transformer.py
Updated rotary position embedding implementation	Modified buffer handling for rotary embeddings Added support for cos/sin separate tensors Updated position embedding application in attention mechanism	megatron/model/rotary_pos_embedding.py megatron/model/language_model.py
Added example configuration for FPDT training	Created new training script for GPT model with FPDT Added configuration parameters for FPDT chunk size and offloading Included sequence parallel size settings	examples_deepspeed/sequence_parallel/ds_pretrain_gpt_6.7B_fpdt_32k.sh

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time. You can also use
  this command to specify where the summary should be inserted.

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey @saforem2 - I've reviewed your changes - here's some feedback:

Overall Comments:

• Consider adding documentation explaining the FPDT optimization approach and when it should be used vs standard execution

Here's what I looked at during the review

• 🟡 General issues: 2 issues found
• 🟢 Security: all looks good
• 🟢 Testing: all looks good
• 🟡 Complexity: 1 issue found
• 🟢 Documentation: all looks good

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[sourcery-ai]

issue (performance): Moving from register_buffer to regular tensor attribute will cause unnecessary tensor recreation

Consider keeping this as a register_buffer to avoid recreating the tensor on every forward pass

[sourcery-ai]

issue: Sequence length divisibility requirements should be validated explicitly

Add explicit validation that sequence length is divisible by required factors to fail fast with a clear error message

[sourcery-ai]

issue (complexity): Consider refactoring the parallel attention implementation to use a cleaner class hierarchy with FPDT as the primary path

The parallel attention implementation has become overly complex with duplicate code paths. Consider refactoring to make FPDT the primary implementation:

class ParallelAttention:
    def __init__(self, config, ...):
        self.use_legacy = not config.use_fpdt
        self.fpdt_attention = FPDTAttention(...) if not self.use_legacy else None
        self.legacy_attention = LegacyAttention(...) if self.use_legacy else None

    def forward(self, hidden_states, attention_mask, ...):
        if self.use_legacy:
            return self.legacy_attention(hidden_states, attention_mask, ...)
        return self.fpdt_attention(hidden_states, attention_mask, ...)

This separates the implementations while maintaining compatibility. The FPDT path removes unnecessary dropout/bias complexity. Consider deprecating the legacy path in future versions.