slurm Multi-machine and multi-GPU training

[yangzhipeng1108]

2: Traceback (most recent call last):
2: File "/workspace/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py", line 64, in main
2: model = MegatronGPTSFTModel.restore_from(cfg.model.restore_from_path, model_cfg, trainer=trainer)
2: File "/opt/NeMo/nemo/collections/nlp/models/nlp_model.py", line 465, in restore_from
2: return super().restore_from(
2: File "/opt/NeMo/nemo/core/classes/modelPT.py", line 464, in restore_from
2: instance = cls._save_restore_connector.restore_from(
2: File "/opt/NeMo/nemo/collections/nlp/parts/nlp_overrides.py", line 1122, in restore_from
2: loaded_params = super().load_config_and_state_dict(
2: File "/opt/NeMo/nemo/core/connectors/save_restore_connector.py", line 178, in load_config_and_state_dict
2: instance = calling_cls.from_config_dict(config=conf, trainer=trainer)
2: File "/opt/NeMo/nemo/core/classes/common.py", line 524, in from_config_dict
2: raise e
2: File "/opt/NeMo/nemo/core/classes/common.py", line 516, in from_config_dict
2: instance = cls(cfg=config, trainer=trainer)
2: File "/opt/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_sft_model.py", line 77, in init
2: super().init(cfg, trainer=trainer)
2: File "/opt/NeMo/nemo/collections/nlp/parts/mixins/nlp_adapter_mixins.py", line 79, in init
2: super().init(*args, **kwargs)
2: File "/opt/NeMo/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py", line 287, in init
2: super().init(cfg, trainer=trainer, no_lm_init=True)
2: File "/opt/NeMo/nemo/collections/nlp/models/language_modeling/megatron_base_model.py", line 175, in init
2: init_world_size = trainer.world_size
2: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1161, in world_size
2: return getattr(self.strategy, "world_size", 1)
2: File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/parallel.py", line 66, in world_size
2: return self.cluster_environment.world_size() if self.cluster_environment is not None else 1
2: File "/usr/local/lib/python3.10/dist-packages/lightning_fabric/plugins/environments/torchelastic.py", line 63, in world_size
2: return int(os.environ["WORLD_SIZE"])
2: File "/usr/lib/python3.10/os.py", line 680, in getitem
2: raise KeyError(key) from None
2: KeyError: 'WORLD_SIZE'

use slurm

#!/bin/bash
#SBATCH --exclusive --nodes=2 --mem=0 --overcommit --ntasks-per-node=8 --time=4:30:00 --job-name=multinode_sft_example

#Load necessary modules and set environment variables
export CUDA_DEVICE_MAX_CONNECTIONS=1

#Set model and training parameters
TRAIN_DS="[/workspace/data/databricks-dolly-15k/training.jsonl]"
VALID_DS="[/workspace/data/databricks-dolly-15k/validation.jsonl]"
TEST_DS="[/workspace/data/databricks-dolly-15k/test.jsonl]"
CKPT="/workspace/model_ckpts/llama2-7b.nemo"

TS=$(date +%s)
OUTPUT_PATH="/workspace/"
RESULTS_DIR="$OUTPUT_PATH/results_${TS}"
CONCAT_SAMPLING_PROBS="[1]"

TP_SIZE=1
PP_SIZE=1
BS=128
MAX_LEN=2048

#The NeMo command to run on each node.
run_cmd="python3 /workspace/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py
trainer.precision=bf16
trainer.devices=$SLURM_NTASKS_PER_NODE
trainer.num_nodes=$SLURM_JOB_NUM_NODES
trainer.val_check_interval=0.1
trainer.max_steps=50
model.restore_from_path=${CKPT}
model.micro_batch_size=1
model.global_batch_size=128
model.tensor_model_parallel_size=${TP_SIZE}
model.pipeline_model_parallel_size=${PP_SIZE}
model.megatron_amp_O2=True
model.activations_checkpoint_granularity=selective
model.activations_checkpoint_method=uniform
model.optim.name=distributed_fused_adam
model.optim.lr=5e-6
model.answer_only_loss=True
model.peft.peft_scheme=none
model.data.train_ds.file_names=${TRAIN_DS}
model.data.validation_ds.file_names=${VALID_DS}
model.data.test_ds.file_names=${TEST_DS}
model.data.train_ds.concat_sampling_probabilities=${CONCAT_SAMPLING_PROBS}
model.data.train_ds.max_seq_length=2048
model.data.validation_ds.max_seq_length=2048
model.data.train_ds.micro_batch_size=1
model.data.train_ds.global_batch_size=128
model.data.validation_ds.micro_batch_size=1
model.data.validation_ds.global_batch_size=128
model.data.test_ds.micro_batch_size=1
model.data.test_ds.global_batch_size=256
model.data.train_ds.num_workers=0
model.data.validation_ds.num_workers=0
model.data.test_ds.num_workers=0
model.data.validation_ds.metric.name=loss
model.data.test_ds.metric.name=loss
exp_manager.create_wandb_logger=False
exp_manager.explicit_log_dir=/workspace/results
exp_manager.resume_if_exists=True
exp_manager.resume_ignore_no_checkpoint=True
exp_manager.create_checkpoint_callback=True
exp_manager.checkpoint_callback_params.monitor=validation_loss
exp_manager.checkpoint_callback_params.save_best_model=False
exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True
++cluster_type=BCP"

#Container
CONT=/root/nvcr.io+nvidia+nemo+24.06.sqsh
CONT_NAME=nemofw-training
CONT_MOUNT=/root/yovole:/workspace,/dev/infiniband:/dev/infiniband

#run on SLURM
srun -l
--ntasks-per-node=8
--container-name="${CONT_NAME}"
--container-image="${CONT}"
--container-mounts="${CONT_MOUNT}"
--container-entrypoint
--no-container-mount-home
bash -c "${run_cmd}"