NeMo Distillation API
State: closed | Pull request: yes
Labels: name: NLP

name: CI

name: Run CICD

# What does this PR do ?

Add Logits-only Distillation ability to NeMo GPT models.

**Collection**: NLP

# Changelog 
- Full-fledged example script containing a distillation-enabled subclass of `MegatronGPTModel` and proper loss function for logit-based knowledge distillation.
- Default config file, docs, and CICD test

# Usage

```bash
STUDENT_CKPT=".../model1.nemo"  # can also be None
TEACHER_CKPT=".../model2.nemo"
TOKENIZER=".../tokenizer.model"
FINAL_SAVE_FILE="final_checkpoint.nemo"
TP=4
PP=1
NPROC=$((TP * PP))
launch_config="torchrun --nproc_per_node=$NPROC"

${launch_config} examples/nlp/language_modeling/megatron_gpt_distillation.py \
    model.restore_from_path=$STUDENT_CKPT \
    model.kd_teacher_restore_from_path=$TEACHER_CKPT \
    model.tensor_model_parallel_size=${TP} \
    model.pipeline_model_parallel_size=${PP} \
    model.micro_batch_size=1 \
    model.tokenizer.model=$TOKENIZER \
    model.nemo_path=$FINAL_SAVE_FILE \
    trainer.precision=bf16 \
    trainer.devices=$NPROC
```

# 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**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [x] Did you write any new necessary tests?
- [x] Did you add or update any necessary documentation?
- [x] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [x] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [x] 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)


@pablo-garay @ko3n1g  I'd like to request a review, thanks!

@pablo-garay is OOTO til Aug 24.

@ko3n1g You're a code owner for same file, so it's redundant.  Could we bypass Pablo so we can merge?

hey @AAnoosheh, I'm still working on a fix for the CI. Unless it's highly business critical I would prefer waiting until our infra is up again such that we fully test this PR

@ChenhanYu How business critical is this atm?

Hi I want to use this code for Distillation.
I was not able to find **hide_teacher_model** function in "examples/nlp/language_modeling/megatron_gpt_distillation.py" and I got error.

Where can I find **hide_teacher_model** function?


Hi, @sangHa0411 

The `hide_teacher_model` is part of the Distillation module in the [Model Optimizer library](https://github.com/NVIDIA/TensorRT-Model-Optimizer) but not officially publicized yet. 

If you have v0.15 of the library installed and use the script as is, it should be present in the model after it gets converted – on [this line](https://github.com/NVIDIA/NeMo/blob/d7d6b99c22d3edaeb0f61b2b1719e4e400f540af/examples/nlp/language_modeling/megatron_gpt_distillation.py#L160)

“Pipeline parallelism not yet supported”，then how to deal with the models like llama 70B? Do we apply the tensor parallelism 
 across multiple nodes?

@lyravv
Llama-70b can't even be tensor-parallel split further than 8 devices as it has GQA with 8 groups.
It would require shard duplication, and I'm not sure if NeMo already supports that or not.

But PP support is underway and hopefully available soon.

Now the TP has much less TFLPs than PP. So PP support is very important and I am very interest in it.

By the way, how to print the `ce_loss` and `kd_loss` before combine them together? I tried to add `print` in /usr/local/lib/python3.10/dist-packages/modelopt/torch/distill/loss_balancers.py at `StaticLossBalancer::forward`, but it seems the logging is disabled and no output. Does anyone has similar issue? Thanks.