Can I use 'torchtitan' for any model architecture?

[jefcoder]

Hi torchtitan team and torchtitan users,

I currently use high-level libraries such as Axolotl and ms-swift for LLM training. While these are convenient and easy to use, they make it difficult to access and modify the full training pipeline end-to-end.

I've found that torchtitan looks like a great choice for my case —it offers many advanced features and performance. However, as far as I understand, it appears to be primarily built around the Llama family of models. I'd like to know more about how well torchtitan supports other model architectures.

Specifically, I'm currently developing a new model by modifying the architecture of Qwen, using custom Python implementation files. Given this use case, is torchtitan a suitable choice for building an end-to-end training pipeline for a custom architecture like this? Or would it be more advisable to rely on Hugging Face Transformers and TRL instead?

Thank you for your time and insight.

[tianyu-l]

Please check https://github.com/pytorch/torchtitan/tree/main/torchtitan/models