Fine-Tuning Llama Model with Large Context and Customized Dataset Using Torchtitan

[Amerehei]

Hi,

I am trying to fine-tune a Llama model with a large context size, and I found that to efficiently shard activations across multiple GPUs, I need to use Torchtitan. Here are some questions related to my setup:

See related issue: meta-llama/llama-recipes#785

1. Custom Dataset Usage
   I created a custom dataset using parquet files and a custom_dataset.py file, which is compatible with llama-recipes. I'm also using the DEFAULT_CHATML_CHAT_TEMPLATE. Could you please provide guidance on how to integrate and use this custom dataset effectively with Torchtitan?

2. Fine-Tuning with Pretrained Model
   Is it possible to fine-tune the model starting from a pretrained checkpoint? If so, are there specific steps or configurations needed to achieve this with Torchtitan?

3. Model Support (Llama-3.2-1B)
   I noticed that Torchtitan currently supports training Llama 3 models (8B, 70B) out of the box. What steps would I need to take if I wanted to train meta-llama/Llama-3.2-1B specifically?

4. Large Context and FSDP Limitation
   I am unable to use FSDP because of the large context sizes I’m working with. Any additional guidance on handling large contexts effectively with Torchtitan would be appreciated.

Thank you for your help!

[aniltrkkn]

For 3.1) You need to implement apply_scaling on top of Llama 3 code

https://github.com/meta-llama/llama-models/blob/main/models/llama3/reference_impl/model.py

Note: I just realized you asked for 3.2, I read it as 3.1. But linked library has 3.2 implementations which are trivial to port.

[gnadathur]

cc: @tianyu-l
cc: @XilunWu on (4). @Amerehei -- for (4) titan now supports FSDP + context parallel for long seq length training

[gnadathur]

cc: @lessw2020

[gnadathur]

cc: @fduwjj on LLaMa3.2 support.

[XilunWu]

Regarding 4)
The PR #592 enables CP in torchtitan. You can change context_parallel_degree (for example 8 for Cp8) in the toml file. See detail in the PR description. CP8 is enough for 128K on H100 and A100. If you still encounter OOM, you can change selective checkpoint to "full" to further reduce peak memory usage.

Let me know if you need more assistance on enabling large sequence input. Glad to help.

[tianyu-l]

1. Custom Dataset Usage

Currently we provide a simple data loading solution in https://github.com/pytorch/torchtitan/blob/main/torchtitan/datasets/hf_datasets.py
Do you think you can adapt the code there for your use case? In principle, as long as your dataset comes in as a map-style Dataset / IterableDataset, adaptation should be straightforward.

2. Fine-Tuning with Pretrained Model

Do you find the script https://github.com/pytorch/torchtitan/blob/main/docs/checkpoint.md to be helpful? Basically torchtitan supports DCP format checkpoint loading. To use it, you'll need to convert your checkpoint into DCP format. (What format is your checkpoint in?)

3. Model Support (Llama-3.2-1B)

The support for Llama 3.2 is ongoing in torchtitan. You'll likely need to wait a bit unless you'd try to integrate it by yourself.

4. Large Context and FSDP Limitation

Please try Context Parallel and give feedbacks!

[Amerehei]

Thank you all for your help, I need to read some of torchtitan source codes.
I'll reply here if I have any question