add a feature to drop checkpoint layers
State: closed | Pull request: yes
Labels: name: NLP

name: CI

name: Run CICD

# What does this PR do ?

1. Adds a feature similar to mcore's --skip-train flag.
3. Adds drop-layers feature support.
4. Adds script which loads the ckpt, trims needed layers, and saves trimmed checkpoint.

**Collection**: [Note which collection this PR will affect]

# Changelog 
- Add specific line by line info of high level changes in this PR.

# Usage
* Trim the model. Please, note that layer indices start from 1. Llama3 8b model example:
```ubuntu
python -m torch.distributed.launch --nproc_per_node=8 \
     NeMo/examples/nlp/language_modeling/megatron_gpt_drop_layers.py \
     --path_to_nemo /path/to/model.nemo \
     --path_to_save /path/to/save/trimmed_model.nemo \
     --tensor_model_parallel_size 8 \
     --pipeline_model_parallel_size 1 \
     --gpus_per_node 8 \
     --drop_layers 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31
```
If you want to save trimmed model in zarr ckpt format (deprecated), add this flag to the command:
```ubuntu
--zarr
```

* Validate the model:
```ubuntu
python NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py \
  --config-path=/path/to/folder/with/model/config \
  --config-name=llama3_8b.yaml \
  trainer.limit_val_batches=number_of_batches_in_validation_step \
  model.restore_from_path=/path/to/trimmed_model.nemo \
  model.skip_train=True \
  model.data.data_impl=mock \
  model.data.data_prefix=[]
```
To use specific dataset instead of mock dataset:
```ubuntu
model.data.data_impl=mmap \
model.data.data_prefix=["path/to/datafile1", "path/to/datafile2"]
```

* Validate Original Model

To validate the original model without specific layers, use the following script:

```ubuntu
  python /NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py \
    --config-path=/path/to/folder/with/model/config \
    --config-name=llama3_8b.yaml \
    trainer.limit_val_batches=<limit_val_batches> \
    model.restore_from_path=/path/to/original_model.nemo \
    model.skip_train=True \
    model.data.data_impl=mock \
    model.data.data_prefix=[] \
    model.drop_layers=[16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31]
```

# GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)


This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

Looks good to me. 

Approving MR as review from Marcin is complete.