DeepSpeedCheckpoint: support custom final ln idx

[nelyahu]

till today only last layer (idx=-1) was considered using FINAL_LAYER_NORM_INDEX which is set to -1.
this PR allows the user to pass custom value for model where this default value does not apply.
see example for usage in HabanaAI/Megatron-DeepSpeed fork repository:
https://github.com/HabanaAI/Megatron-DeepSpeed/blob/c9feb8cacabc6dd4da4266cff08db555a21122e2/tools/verify_checkpoint_non_tp_consistency.py#L296

[nelyahu]

@loadams can you please re-run "nv-torch-latest-v100" validation ? i think it failed on set-up issue

[loadams]

@loadams can you please re-run "nv-torch-latest-v100" validation ? i think it failed on set-up issue

Yes, re-running now and will work on getting this merged.

[jinyouzhi]

Sorry disturb you here, Could you explain why FINAL_LAYER_NORM_INDEX set to -2 not -1 for LLaMA? @nelyahu Thanks.

[nelyahu]

@jinyouzhi The previous code assumed that model is built of embedding + transformer layers + projection layer
Therefore, in order to get all transformer layer it took layers[1:-1] (i.e. excluding the embedding and the projection layers).
However, llama has an RMSNorm layer before the projection layer, so the above assumption is incorrect - and need to exclude also the 2nd to last as well (i.e. embedding, RMSNorm and projection).

Also this approach of fetching layers based on their indices is not a good practice, and need to invest more efforts in doing it based on the layer type so it will generic and robust.

[jinyouzhi]

@nelyahu got it. Thank you very much!