mirror of
https://github.com/hanzo-ml/trainer.git
synced 2026-09-24 00:31:34 +00:00
Yash AgarwalandGitHub
36f5569c16
* feat: add Megatron-Core GPT Tensor Parallelism example notebook Signed-off-by: XploY04 <[email protected]> * added megatron notebook to the e2e gpu test Signed-off-by: XploY04 <[email protected]> * feat: Parameterize Megatron core GPT notebook's tensor parallelism and GPU count, and update the e2e test workflow to pass these parameters. Signed-off-by: XploY04 <[email protected]> * change the number of gpus to 2 Signed-off-by: XploY04 <[email protected]> * docs: update Megatron GPT tensor parallelism example notebook. Signed-off-by: XploY04 <[email protected]> * fix: correct minor typos and improve code readability in Megatron-Core GPT notebook Signed-off-by: XploY04 <[email protected]> * fix: update tensor model parallel size retrieval and improve code clarity in Megatron-Core GPT notebook Signed-off-by: XploY04 <[email protected]> * fix: correct wording in training function description for clarity in Megatron-Core GPT notebook Signed-off-by: XploY04 <[email protected]> * feat: add verification step for TrainJob completion in Megatron-Core GPT notebook Signed-off-by: XploY04 <[email protected]> * fix: use multi-node setup for Megatron TP to work with GPU time-slicing GPU time-slicing (replicas=2) advertises 2 GPUs to Kubernetes but exposes only 1 CUDA device inside the container. This caused torchrun to launch 1 worker (WORLD_SIZE=1), failing Megatron's TP requirement of world_size >= 2. Switch from 1 node with 2 GPUs to 2 nodes with 1 GPU each. This creates 2 pods, each getting 1 time-sliced GPU, giving WORLD_SIZE=2 for tensor parallelism across pods. Signed-off-by: XploY04 <[email protected]> * fix: use PyTorch devel image for Megatron compile_helpers() support Megatron-Core's compile_helpers() requires `make` and `gcc` to build C dataset helpers. The default runtime image doesn't include build tools. Switch to the devel image which has them pre-installed. Signed-off-by: XploY04 <[email protected]> * fix: install megatron-core from git for compile_helpers() support The pip package for megatron-core doesn't include the Makefile and C source files needed by compile_helpers(). Install from the git repository instead, which includes all source files for compiling the C++ dataset helpers required by MockGPTDataset. Signed-off-by: XploY04 <[email protected]> * fix: install megatron-core from zip archive instead of git The devel image doesn't have git installed, so pip install from git+https fails. Use the GitHub zip archive URL instead, which pip can download and extract directly without needing git. Signed-off-by: XploY04 <[email protected]> * fix: download Megatron C++ source files for compile_helpers() The megatron-core pip package excludes the Makefile and helpers.cpp needed by compile_helpers(). Download these two files from GitHub at runtime before compilation. Also add pybind11 as an explicit dependency since it's needed by the Makefile but only included as a build-time dependency in the pip package. Signed-off-by: XploY04 <[email protected]> * fix: use correct null tokenizer library name for pip megatron-core The pip release of megatron-core uses "null" as the tokenizer library name, while the GitHub main branch uses "null-text". Change to "null" to match the pip version. Signed-off-by: XploY04 <[email protected]> * fix: increase wait_for_job_status timeout to 1800s The megatron training job takes longer than 900s due to pip install, C++ compilation, and multi-node training setup on time-sliced GPUs. Increase timeout to 1800s to match PAPERMILL_TIMEOUT. Signed-off-by: XploY04 <[email protected]> * fix: use runtime image with apt-get for build tools instead of devel image The devel image (3.5GB) is too slow to pull, causing PAPERMILL_TIMEOUT. Switch back to the runtime image and install only make and g++ via apt-get (~50MB). This provides everything compile_helpers() needs without the overhead of pulling the full devel image. Signed-off-by: XploY04 <[email protected]> * fix: disable NCCL shared memory to avoid /dev/shm limit in Kubernetes Kubernetes pods get 64MB /dev/shm by default. NCCL needs more for its communication buffers, causing checkpoint save to fail with "Error while creating shared memory segment". Set NCCL_SHM_DISABLE=1 to use network sockets instead. Also reduce wait_for_job_status timeout to 60s since training completes quickly and the main wait happens in get_job_logs. Signed-off-by: XploY04 <[email protected]> * debug: add NCCL_DEBUG=INFO to diagnose /dev/shm failure on node-1 Signed-off-by: XploY04 <[email protected]> * fix: mount /dev/shm as emptyDir to fix NCCL shared memory exhaustion NCCL proxy service allocates ~33MB per communicator in /dev/shm. The default Kubernetes /dev/shm is 64MB (Docker default), which is insufficient for workloads that create multiple NCCL communicators. Patch ClusterTrainingRuntimes in GPU e2e setup to mount /dev/shm as a Memory-backed emptyDir, and remove unnecessary NCCL workaround env vars from the Megatron notebook. Signed-off-by: XploY04 <[email protected]> * fix: remove dist_checkpointing from Megatron notebook Megatron-Core dist_checkpointing uses multiprocessing.spawn internally to create a Manager queue for async writes. The Kubeflow SDK generates training scripts without an if __name__ == '__main__' guard, so the spawned child re-imports the script and re-executes the training function, causing a RuntimeError. Remove the checkpoint step since training (Steps 1-5) is the core TP demonstration. Also remove unused imports and fix the GPU prerequisites text. Signed-off-by: XploY04 <[email protected]> --------- Signed-off-by: XploY04 <[email protected]>