Add Nemotron model via PAXML to CloudAI + optimization for large GPU runs

[srivatsankrishnan]

Summary

This is an umbrella PR for supporting large GPU runs in CW. Though this PR was originally created for adding just Nemotron, there were bunch of feature request and changes for having resilient/best known method s that were needed to scale beyond 1k GPUs. As of today, we have successful runs upto 2K GPUs. The following features were added to via this PR

• Nemotron models via PAXML.
• XLA flags specific to perf/profile stages
• Enable/Disable PGLE flow through a flag.
• Start the container and launch the training through container name instead of sqsh file
• Jax training does not need the mpi flag in srun. For NCCL pre-test it is okay
• Perfomance fix for profiling stage in CloudAI. CloudAI jobs were ~4x slower (measured at 1k gpus jobs) compared to standalone bash scripts.
• Nsys profiling only on rank-0 jobs (was slowing up jobs at 1k+ gpu scale)
• Remove the nsys trace to sqlite converter (not needed as per Haixin)

Test Plan

• CI/CD should pass. The existing unit test should support the new architecture.
  More details will be updated later.

• Test on internal systems to make sure it trains GPT/Grok/Nemotron correctly meeting existing performance targets
  More details will be updated later. We were able to successfully scale to 2K gpu runs for Grok-1 via Paxml.

$ python cloudaix.py --mode run --system-config conf/internal/jax_toolbox/system/israel_1.toml --test-templates-dir conf/internal/jax_toolbox/test_template/ --tests-dir conf/internal/jax_toolbox/test --test-scenario conf/internal/jax_toolbox/test_scenario/nemo_proxy/test_scenario_israel_1.toml
[INFO] System configuration file: conf/internal/jax_toolbox/system/israel_1.toml
[INFO] Test templates directory: conf/internal/jax_toolbox/test_template
[INFO] Tests directory: conf/internal/jax_toolbox/test
[INFO] Test scenario file: conf/internal/jax_toolbox/test_scenario/nemo_proxy/test_scenario_israel_1.toml
[INFO] Output directory: None
[INFO] System Name: Israel-1
[INFO] Scheduler: slurm
[INFO] Test Scenario Name: jax_toolbox
[INFO] Checking if test templates are installed.
[INFO] Test Scenario: jax_toolbox

Section Name: Tests.1
  Test Name: sanity_fp8_nemo340b_1
  Description: sanity-fp8-nemo340b-1
  No dependencies
[INFO] Initializing Runner
[INFO] Creating SlurmRunner
[INFO] Starting test scenario execution.
[INFO] Starting test: Tests.1
[INFO] Running test: Tests.1
[INFO] Executing command for test Tests.1: sbatch /results/jax_toolbox_2024-08-17_00-23-43/Tests.1/0/cloudai_sbatch_script.sh

[amaslenn]

Looks like this PR depends on #170, it is a bit hard to understand for is uniq for this PR itself. Especially in the context of testing. Please try to keep the coverage at least on the same level as it is.

[srivatsankrishnan]

Looks like this PR depends on #170, it is a bit hard to understand for is uniq for this PR itself. Especially in the context of testing. Please try to keep the coverage at least on the same level as it is.

This PR adds lot of features required for running 1k-4k gpu jobs. Updated the the PR summary to reflect this.