how to adapt my new model?

[belle9217]

[felipemello1]

Hi @belle9217 , can you give more details of what you are trying to do?

[belle9217]

I have a model with a similar architecture to the LLaMA model, which has been modified based on LLaMA by altering the multi-head attention mechanism for QKV. now I want to use FSDP(deepspeed zero3) and torch.compile,and I found torchtune has the two functionalities.so I want to add my model to the torchtune framework. since I see torchtune has support Qwen、llama、phi, how can I adjust my modeling.py to use torchtune? is there any API or Interface？

[calvinpelletier]

Hi @belle9217 , you'll want to create a model builder, which is a function that returns an torch.nn.Module like this. Then point to the function in your model config like this.

Choose a checkpointer (which is responsible for model loading/saving) in your config like this. You can find our checkpointers here. You may need to write a custom checkpointer if your weights are in a different format. The load_checkpoint function is how torchtune will load your weights.

If you want to use LoRA, here's an example builder and an example recipe. We provide LoRA modules and utilities here.

Let me know if you have any questions!