Mitigation to HuggingFace Trainer

[huyiwen]

A lot of people in the community use HuggingFace Trainer for training, but sometimes it’s not flexible enough or missing certain features (native tp/pp/ep etc). Migrating to Megatron-LM comes with a steep learning curve, and while TorchTitan is lighter, it still takes some effort to learn and doesn’t fully support features like Flash Attention and Liger Kernel yet (correct me if I’m wrong).

One way to make TorchTitan more accessible could be allowing some of its features to work with existing HuggingFace Trainer code with just minor tweaks—like different parallelisms. That way, more users might give it a try even if it doesn’t fully support all training features yet.

[tianyu-l]

@huyiwen Thanks for asking. This is something we are considering.

One way to make TorchTitan more accessible could be allowing some of its features to work with existing HuggingFace Trainer code with just minor tweaks—like different parallelisms.

I wonder if you could provide a more concrete list of requirements / to-do items in order to integrate with "HF Trainer"?

[huyiwen]

I’m currently working on an MoE model and looking to implement expert parallelism. Writing EP/EP+TP/EP+DP from scratch with torch distributed communication is pretty challenging, especially if I want good training speed. That’s why using the parallelism provided by DTensor and torchtitan seems like a solid option.

Beyond EP, it might also be a good idea to enable people to try out PP and TP without relying on Megatron-lm. This way, more people could train larger models with limited resources.

[tianyu-l]

For the most part, torchtitan doesn't implement the parallelisms themselves -- the parallelisms core code are in pytorch.

Is your request that "pytorch should make its parallelisms easily usable in HF Trainer"? In fact AFAIK HF already integrates several of pytorch parallelisms. cc: @kwen2501

Can you be more specific on how would you hope torchtitan to adapt?

[yzhangcs]

Hi, @huyiwen I think the smoothest way for migration is to define hf style models and use torchtitan for training with 4d parallel. You may also be interested in https://github.com/fla-org/flame.git, which imo is exactly what you are looking for.

And I also believe that liger kernels might not be a question. After several days of practice, I found that torch.compile+4d would be fast enough. There are a lot of questions remaining unresolved when you try to combine fused kernels with 4d and compile.

[huyiwen]

Thanks, @tianyu-l and @yzhangcs! This is more of a general-purpose idea. If the torchtitan team could test existing and upcoming features (like parallelism) with a basic HF Trainer, it might speed things up for the community to adopt these feature, and make it easier for HF to integrate them too. It feels kind of similar to what the liger-kernel is doing, like having day-1 compatibility with HF Trainer.

[kwen2501]

Hi @yzhangcs thank you for your comments.

I think the smoothest way for migration is to define hf style models and use torchtitan for training with 4d parallel

Dumb q: would a HF-style model definition suffice to run with HF Trainer? Did HF document the style requirement somehwere?

You may also be interested in https://github.com/fla-org/flame.git

Thank you for the pointer, can you please point us to the exact feature / API of Flame you refer to? Thanks!

[yzhangcs]

Hi, @kwen2501

Dumb q: would HF-style model definition enable composability with HF Trainer? Did HF document the style requirement somehwere?

HF-style models (e.g., AutoModelForCausalLM) also inherit from nn.Module, making them inherently decoupled from the HF Trainer. This allows seamless integration with frameworks like torchtitan, following similar patterns to torchtitan's existing model recipes.

Implementation Nuances:

• Initialization: HF's complex weight initialization logic can be streamlined. In flame, we reset all init flags on meta device (code), then re-initialize with proper settings on their respective devices after FSDP (code).
• Parallelisms: We adapted torchtitan's plan registration logic (see implementation), and module naming conventions are critical.

other things are exactly the same!

Did HF document the style requirement somehwere?

I just skimmed through the docs.
The fla models in flame closely follow HF's modeling_llama architecture as a reference.