How to add supervised finetuning mask in torchtitan?

[kailashg26]

How do I implement supervised fine-tuning (SFT) masking in TorchTitan for posttraining using a synthetic dataset?

[kailashg26]

@tianyu-l , is there a PR on going for supporting SFT mask in torchtitan? or is this a different issue?
#1321

[tianyu-l]

There is no ongoing effort right now in torchtitan. I think the main work here is to support

1. a new dataset example
2. masking in loss computation

IIUC 2 should be shared with VLM training as well.
cc @wwwjn @lkhphuc

[kailashg26]

@tianyu-l Thanks for quick response. If I understand correctly, the above PR is using C4 dataset to finetune the pretrained llama 3 8B model? C4 is basically the same dataset used for pre-training as well right? Please lmk if I'm missing something :)
Also, Could you please share some insights into how to do SFT in torchtitan? Would be really helpful.

[kailashg26]

@wwwjn @lkhphuc @tianyu-l Could you please provide some insights into this?

[tianyu-l]

I think the main missing piece is data loading for SFT.

For examples of training on a fine-tuning dataset, you could refer to how torchtune handles SFT data loading, e.g. https://github.com/meta-pytorch/torchtune/blob/main/torchtune/datasets/_sft.py

If we could find bandwidth, we may consider adding an example in torchtitan.

[kailashg26]

I can actually try adding one for torchtitan. But I'm more familiar with torchtune. So, if you could please help with steps or something like that I can start working on it :) Thanks @tianyu-l

[tianyu-l]

Oh, is it hard to switch torchtitan dataloader to a torchtune one? Curious what difficulties you hit?

[kailashg26]

My bad, should have explained better. Since torchtune has stopped dev tasks. I was curious to try finetuning on torchtitan. So, I was wondering if we can do finetuning on torchtitan. From my understanding in order to try SFT on torchtitan, I need to change the data loader, loss masking, and the dataset should be prompt-response right? Please let me know if I'm missing any other things.

[tianyu-l]

change the data loader, loss masking, and the dataset should be prompt-response right?

That sounds right to me.

Another piece which may be involved is tokenizer, as you might need some special tokens for specific tasks, e.g. Q&A, reasoning.