AutoTP for training skips lm_head ?

[3outeille]

I noticed that when using AutoTP for training purpose, the lm_head is skipped. It seems to be only working for inference but not training. Any ideas why ? cc:@inkcherry

[inkcherry]

hi, @3outeille , you can try this #7181 , currently, this has not been turned into a configurable setting yet. btw (Compared to using configuration, I believe leveraging Hugging Face's TP plan would be a better approach, It will take a bit of time as I’ve been occupied with other things lately) contributions are welcome.

[3outeille]

I believe leveraging Hugging Face's TP plan would be a better approach, It will take a bit of time as I’ve been occupied with other things lately) contributions are welcome.

I am actually taking a look at it and one painpoint is that Huggingface's TP plan leverage Dtensor while deepspeed doesnt. This means a lot of patching around deepspeed function to convert dtensor to normal torch tensor using to_local() but not sure if it is the best approach. Do you have a better recommendation ?