Vote on new features!

[tianyu-l]

Hi torchtitanists,

Thank you for your interests in torchtitan!

Please upvote on what features you would like to see next, and add one if it's not already there. We'll try to prioritize on the most requested features.

[tianyu-l]

#625 multimodal models (e.g. Llama 3.2) (work in progress)

[tianyu-l]

#184 MoE and Expert Parallel (work in progress)

[mayank31398]

scattermoe might be useful?
https://arxiv.org/abs/2403.08245

[SeunghyunSEO]

to my best knowledge, scattermoe is compatible with FSDP (maybe 1) and IBM granite is trained with it successfully, but it does not support EP and FSDP2, right?
i want torch native EP haha

[mayank31398]

correct, we train Granite with scattermoe.
It works with FSDP-2 + TP as well.
this is the TP implementation: https://github.com/IBM/dolomite-engine/blob/main/dolomite_engine/hf_models/models/moe_dolomite_TP/moe_TP/scatter.py
this is the non-TP implementation: https://github.com/IBM/dolomite-engine/blob/main/dolomite_engine/hf_models/models/moe_dolomite/moe/scatter.py

[SeunghyunSEO]

correct, we train Granite with scattermoe.

It works with FSDP-2 + TP as well.

this is the TP implementation: https://github.com/IBM/dolomite-engine/blob/main/dolomite_engine/hf_models/models/moe_dolomite_TP/moe_TP/scatter.py

this is the non-TP implementation: https://github.com/IBM/dolomite-engine/blob/main/dolomite_engine/hf_models/models/moe_dolomite/moe/scatter.py

oh, thank you so much for sharing.
ill try scatter moe with torchtitan or lingua :)

[casper-hansen]

This fork has some code for "InternLM3 MoE" that implements a dropless MoE. It's not a ScatterMoE implementation, but does use Grouped GEMM.
https://github.com/zigzagcai/mytorchtitan

[tianyu-l]

state space models (e.g. Mamba)

[tianyu-l]

diffusion models (e.g. DiT)

[gil2rok]

this!

[tangjiasheng]

maybe this need to open a new repo.

[ad8e]

What about varlen ring attention? There are a few variants which are not stable enough for production use; its engineering/testing obstacles make it suitable for the strengths of the PyTorch folks.

[casper-hansen]

Evaluations implementation with support for 4D parallelism

[tangjiasheng]

I'd like to see a Deepspeed-Ulysses-Style sequence parallel used with DTensor.

[casper-hansen]

Qwen 2.5 support. Features needed on top of llama.py:

• safetensors to DCP
• attention bias on QKV
• YaRN (implemented in qwen2.py in HF transformers)

[conceptofmind]

Definitely should add YaRN

[dz1iang]

Parallel training is non-intrusive to model implementation.

[tianyu-l]

Can you elaborate more on this request? I believe our current parallelisms have largely avoided intrusive changes to model implementations.

[dz1iang]

Can you elaborate more on this request? I believe our current parallelisms have largely avoided intrusive changes to model implementations.

I have an idea to modify parallelize_llama.py into parallelize_model.py, allowing for quick parallel training of HuggingFace models.

[samsja]

Would be nice to have a native (and tested) transformer export script as well. There is some trick to do with the complex number rope implementation in the conversion. We have a script to do this that should be compatible with the torchtitan implementation as it was originally copied from it : https://github.com/PrimeIntellect-ai/prime/blob/main/scripts/export_dcp.py

[samsja]

Would be nice as well to have an implementation of sequence packing / document masking using flex attention like in torchtune

[OrenLeung]

It seems like the biggest thing missing from torchtitan that most production gaint LLM training runs have is fault tolerance training

I.e. can quickly recover from a failure by removing just 1 data replicas instead of stopping the whole workload and restarting from ckpt

[cailun01]

• Training data preprocessing examples.
• Multimodal model training.

[eminorhan]

• A complete example script demonstrating end-to-end evaluation of saved checkpoints (e.g. using torchtune and lm-eval).

[casper-hansen]

I already put this as a comment and I agree it would be nice if we can use something standardized like you mentioned.
#693 (comment)

[eminorhan]

• Handling more complicated datasets (e.g. interleaved datasets with given mixing probabilities, support for shuffling, etc.).