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

[tianyu-l]

#184 MoE and Expert Parallel

[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

[casper-hansen]

I'm linking this issue here as it pertains to the overall features for MoE.
DeepSeek V3 Support: #760

[tianyu-l]

state space models (e.g. Mamba)

[tianyu-l]

diffusion models

[giladturok]

this!

[tangjiasheng]

maybe this need to open a new repo.

[tianyu-l]

added here https://github.com/pytorch/torchtitan/tree/main/torchtitan/experiments/flux

[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.

[janEbert]

Should now be implemented by this, right?
https://discuss.pytorch.org/t/distributed-w-torchtitan-breaking-barriers-training-long-context-llms-with-1m-sequence-length-in-pytorch-using-context-parallel/215082

[tianyu-l]

added by #1001
Context Parallel support is ongoing

[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

[tianyu-l]

added by #1001

[samsja]

thanks !!!

[functionstackx]

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

[zhangmozhe]

This is indeed useful

[tianyu-l]

added in #834

[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.).

[samsja]

if it can help I have a simple interleaved / mixing probability dataset implementation here : https://github.com/PrimeIntellect-ai/prime/blob/main/src/zeroband/data.py

[zhangmozhe]

Reinforcement learning with verifiable rewards as used in TULU3 and DeepSeek R1, with asynchronous RL training

[SeunghyunSEO]

bro torchtitan is not for post-training.
torchtitan aims to showcase the cutting-edge large scale distributed training features of the latest pytorch in a clean, minimal codebase.
you should request RLVR feature at torchtune

[janEbert]

Post-training conversion to HuggingFace.

This is quite the beast due to the various options that TorchTitan offers, and not all of them being part of the model code (e.g., FP8 conversion, CP requiring an external context manager, etc.). All that combined with the maintenance aspects of future changes.
My current method for achieving at least evaluation support with an HF-compliant forward front-end is a client-server implementation, in which a distributed TorchTitan model (which is entirely and only dependent on TorchTitan code), is queried from a non-distributed HF model. However, this doesn't handle any of the post-training use cases and isn't really in line with HF's APIs; however, this implementation minimizes maintenance due to adapting most TorchTitan changes automatically.

[samsja]

hey @tianyu-l

Yes eval via lm_eval would work as well.

Also HF implementation is the default implementation for llama code, and any model trained with torchtitan will likely be released and eval via hf impl alter on. Therefore, I believe that TorchTitan should maintain an export function to hf safetensors format. It would also make model train with torchtitan compatible with a broader set of open source tool including vllm, sglang etcc ...

Also wouldn't HF model put restrictions on the sets of parallelisms a model can run -- e.g. can HF model run Pipeline Parallel / Context Parallel?

I think that for eval TP / DP is enough at least up to 70B model on 8H100

[janEbert]

The reason for me to have HF compatibility is because there are several other frameworks, such as the excellent EuroEval suite, which potentially include evaluation tasks that are not available in lm_eval. :)
It is simply a more general API for the world of evaluation, but also enables additional features for free, like chat deployment.

[tianyu-l]

@samsja @janEbert
It sounds that instead of building eval directly in torchtitan, converting to HF format and doing downstream work with HF model definition is more preferable. If that's the case we already have DCP -> .pt -> HF (#758, #420); so maybe we should support direct DCP -> HF conversion?

This is in contrast to the request #693 (comment)
cc @casper-hansen @eminorhan any opinions?

How about in-training eval requested here #883 (comment)
Do people think this feature is useful or not?

[casper-hansen]

I think evals during training is highly valuable. Keep the format in torchtitan for all the parallelism available.

[eminorhan]

My current eval pipeline is dcp -> .pt and then torchtune of .pt (with lm-eval). That's just two lines total, so I haven't really felt a pressing need for in-training evals yet but I can see that it would be useful to utilize the parallelisms already set up during training to make evals more efficient, and in fact this may be necessary for very large models (but I haven't been there yet personally).

On the other hand, I feel like this feature (parallelized eval based on dcp checkpoint) should really be implemented in torchtune or in HF rather than in torchtitan though.