aws-neuron / neuronx-distributed

MIT No Attribution
30 stars 5 forks source link

How can I freeze intermiadiate layers in a model ? #16

Closed shamanez closed 1 month ago

shamanez commented 4 months ago

I am following the llama2 pre-training code.

https://github.com/aws-neuron/neuronx-distributed/blob/main/examples/training/llama2/tp_zero1_llama2_7b_hf_pretrain/tp_zero1_llama2_7b_hf_pretrain.py

I do not understand how to freeze some of the parameters.

aws-rhsoln commented 2 months ago

Thank you for reporting the issue. For freezing weights, you can follow the same technique that one would follow for other torch models, need to set the requires_grad attr to false for the params that needs to be frozen.

aws-rhsoln commented 1 month ago

closing this issue, please re-open if the issue persists.