Support NeMo MegatronGPTModel

[athitten]

🚀 Feature

Support NeMo's LLM MegatronGPTModel

Initial examine:
Found 44 distinct operations, of which 33 (75%) are supported

Motivation

Pitch

Work items

• Implement TensorBase.div_ #340
• Implement TensorBase.__setitem__ #341
• Implement torch.device  #327
• Implement TensorBase.clone #342
• Impement TensorBase.view_as #328
• Implement _set_grad_enabled of torch._C #332
• Implement TensorBase.is_cuda #329
• Implement TensorBase.is_complex #339
• Implement _VariableFunctionsClass.empty of torch #331
• Implement _VariableFunctionsClass.max of torch #338
• Implement _VariableFunctionsClass.baddbmm of torch #326
• Unknown attribute _base inside Megatron core #753

Running the model

Required data

Download the data tarball and extract it into the root of your nemo clone.

NeMo installation

To keep the whole thunder team on the same NeMo revisions, and to prevent having a bunch of "modify file to call thunder.jit()" instructions, we temporarily maintain our own branch of thunder. You can grab it by cloning https://github.com/tfogal/NeMo.git. Make sure you have checked out the tfogal/thunder-nemo branch.

To install NeMo, run python3 -m pip install -e . from the root of the checked-out directory.

Running the network

TMPDIR=./foo-mgpt-train
rm -fr ${TMPDIR}; mkdir -p ${TMPDIR}
HYDRA_FULL_ERROR=1 \
THUNDER_ANNOTATE_TRACES=1 \
NEMO_THUNDER_MEGATRON_GPT=1 \
python3 \
  examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py \
  trainer.devices=1 \
  trainer.num_nodes=1 \
  trainer.precision=32 \
  trainer.max_steps=4 \
  trainer.val_check_interval=4 \
  trainer.enable_checkpointing=False \
  +trainer.limit_val_batches=2 \
  +trainer.limit_test_batches=2 \
  exp_manager.checkpoint_callback_params.save_best_model=False \
  exp_manager.exp_dir=examples/nlp/language_modeling/gpt_sft_results \
  model.peft.peft_scheme=none \
  model.optim.name=distributed_fused_adam \
  model.restore_from_path=./data/nlp/megatron_gpt/starcoder-ci-nemo/megatron_starcoder_tp1_pp1.nemo \
  model.tensor_model_parallel_size=1 \
  model.pipeline_model_parallel_size=1 \
  model.data.train_ds.file_names=[./data/nlp/megatron_sft/quarel.jsonl] \
  model.data.train_ds.num_workers=0 \
  model.data.test_ds.file_names=[./data/nlp/megatron_sft/quarel.jsonl] \
  model.data.validation_ds.num_workers=0 \
  model.data.validation_ds.file_names=[./data/nlp/megatron_sft/quarel.jsonl] \
  model.data.test_ds.num_workers=0 \
  model.data.train_ds.concat_sampling_probabilities=[1.0]

The above command was pulled from NeMo's CI testing for this model and lightly modified to account for other directories and set appropriate environment variables.

cc @tfogal