Skip to content

JAX/MaxText multi-node RCCL overlay #228

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
i-kosarev merged 10 commits into from
Sep 8, 2026
Merged

JAX/MaxText multi-node RCCL overlay #228

Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
ce766ce
feat(primus-maxtext): image that builds MaxText against any RCCL commit
i-kosarev Sep 3, 2026
58c953a
feat(jax-maxtext): report median seconds_per_step alongside throughput
i-kosarev Sep 3, 2026
d81cabf
docs(skill): record the jax-maxtext multi-node findings in mad-slurm-…
i-kosarev Sep 3, 2026
921565f
fix(jax-maxtext): strip --config_path unconditionally
i-kosarev Sep 6, 2026
b1dc0c9
fix(jax-maxtext): keep each replaced librccl's RUNPATH, and gate on ldd
i-kosarev Sep 6, 2026
1942be8
fix(jax-maxtext): compare unresolved soname sets, not counts, in the …
i-kosarev Sep 6, 2026
cbe4d62
fix(skill): pin the RCCL build in the manifest template, and add a 40…
i-kosarev Sep 6, 2026
eebec0e
docs(jax-maxtext): state that seconds_per_step is written for every run
i-kosarev Sep 6, 2026
d6eb908
fix(skill): set XLA_GPU_AUTOTUNE_LEVEL=0 in the maxtext templates
i-kosarev Sep 6, 2026
914f150
fix(jax-maxtext): clear the RUNPATH when the destination had none
i-kosarev Sep 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 124 additions & 0 deletions .claude/skills/mad-slurm-multinode/assets/manifests/primus_maxtext_llama3-8b.template.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,124 @@
{
"built_images": {
"rocm-primus-maxtext-llama3-8b-overlay": {
"model": "jax-maxtext/maxtext_MI355X_llama3_8B-fp8-pretrain",
Comment thread
i-kosarev marked this conversation as resolved.
Resolved
"docker_image": "rocm/jax-training:maxtext-v26.6-<FILL_RCCL_TAG e.g. rccl-7.2.1 or rccl-develop-<sha7>>",
"dockerfile": "docker/primus_maxtext_rccl_overlay.ubuntu.amd.Dockerfile",
"base_docker": "rocm/jax-training:maxtext-v26.6",
"build_duration": 0,
"local_image": true,
"registry_image": null,
"registry": null,
"gpu_vendor": "AMD"
}
},
"built_models": {
"rocm-primus-maxtext-llama3-8b-overlay": {
"name": "jax-maxtext/maxtext_MI355X_llama3_8B-fp8-pretrain",
"url": "",
"dockerfile": "docker/primus_maxtext_rccl_overlay.ubuntu.amd.Dockerfile",
"scripts": "scripts/jax-maxtext/run.sh",
"n_gpus": "-1",
"owner": "mad.support@amd.com",
"training_precision": "fp8",
"multiple_results": "primus_perf_output.csv",
"tags": [
"maxtext",
"jax",
"MI355X",
"llama3_8B-fp8-pretrain",
"fp8",
"training"
],
"timeout": -1,
"args": "--config_path examples/maxtext/configs/MI355X/llama3_8B-fp8-pretrain.yaml",
"additional_docker_run_options": "--privileged --group-add render --shm-size 64G --device=/dev/infiniband --cap-add IPC_LOCK --ulimit memlock=-1 --ulimit nofile=1048576 -v /sys:/sys:ro -v /run/udev:/run/udev:ro"
}
},
"context": {
"docker_env_vars": {
"NCCL_DEBUG": "INFO",
"NCCL_DEBUG_SUBSYS": "INIT,NET",
"NCCL_IB_DISABLE": "0",
"NCCL_NET": "IB",
"NCCL_IB_HCA": "<FILL_NCCL_IB_HCA data-plane HCA list, see references/cluster-types.md>",
"NCCL_IB_GID_INDEX": "<FILL_GID RoCEv2 GID index, see references/cluster-types.md>",
"NCCL_SOCKET_IFNAME": "<FILL_IFNAME management iface, see references/cluster-types.md>",
"GLOO_SOCKET_IFNAME": "<FILL_IFNAME management iface, see references/cluster-types.md>",
"RDMAV_DRIVERS": "<FILL_DRIVER see references/cluster-types.md>",
"IBV_DRIVERS": "<FILL_DRIVER see references/cluster-types.md>",
"RCCL_AINIC_ROCE": "<FILL_RCCL_AINIC_ROCE per archetype, see references/cluster-types.md; delete this key if not needed>",
"NCCL_NET_PLUGIN": "none",
"IBV_SHOW_WARNINGS": "1"
},
"docker_mounts": {
"/dev/infiniband": "/dev/infiniband"
},
"docker_build_arg": {
"BASE_DOCKER": "rocm/jax-training:maxtext-v26.6",
"RCCL_COMMIT": "<FILL_RCCL_COMMIT the commit to build; leave empty to track RCCL_BRANCH>",
"RCCL_BRANCH": "develop",
"RCCL_CMAKE_OPTIONS": "<FILL_RCCL_CMAKE_OPTIONS e.g. -DFAULT_INJECTION=ON; empty for a plain build>",
"BUILD_GPU_TARGETS": "<FILL_GPU_TARGETS gfx950 for MI350X/MI355X, gfx942 for MI300X/MI325X>",
"RDMA_CORE_VERSION": "<FILL_RDMA_CORE_VERSION 63.0 on Broadcom Thor2; leave empty to keep the base image's rdma-core>"
},
"gpu_vendor": "AMD",
"guest_os": "UBUNTU",
"docker_gpus": "0,1,2,3,4,5,6,7"
},
"credentials_required": [],
"summary": {
"successful_builds": [],
"failed_builds": [],
"total_build_time": 0,
"successful_pushes": [],
"failed_pushes": []
},
"deployment_config": {
"target": "slurm",
"slurm": {
"partition": "<FILL_PARTITION>",
"account": "<FILL_ACCOUNT or remove if cluster has none>",
"qos": "<FILL_QOS>",
"nodes": 2,
"gpus_per_node": 8,
"time": "12:00:00",
"output_dir": "./slurm_output",
"exclusive": true,
"network_interface": "<FILL_IFNAME management iface, see references/cluster-types.md>"
},
"distributed": {
"launcher": "primus",
"backend": "nccl",
"port": 29500,
"nnodes": 2,
"nproc_per_node": 8,
"primus": {
"config_path": "examples/maxtext/configs/MI355X/llama3_8B-fp8-pretrain.yaml",
"backend": "MaxText"
}
},
"env_vars": {
"HF_HOME": "${HF_HOME}",
"XDG_CACHE_HOME": "${XDG_CACHE_HOME}",
"PIP_CACHE_DIR": "${PIP_CACHE_DIR}",
"MAD_DATAHOME": "${MAD_DATAHOME}",
"NCCL_DEBUG": "INFO",
"NCCL_DEBUG_SUBSYS": "INIT,NET",
"NCCL_IB_DISABLE": "0",
"NCCL_NET": "IB",
"NCCL_SOCKET_IFNAME": "<FILL_IFNAME management iface, see references/cluster-types.md>",
"GLOO_SOCKET_IFNAME": "<FILL_IFNAME management iface, see references/cluster-types.md>",
"NCCL_IB_GID_INDEX": "<FILL_GID RoCEv2 GID index, see references/cluster-types.md>",
"NCCL_IB_HCA": "<FILL_NCCL_IB_HCA same value as context.docker_env_vars>",
"RDMAV_DRIVERS": "<FILL_DRIVER see references/cluster-types.md>",
"IBV_DRIVERS": "<FILL_DRIVER see references/cluster-types.md>",
"RCCL_AINIC_ROCE": "<FILL_RCCL_AINIC_ROCE per archetype, see references/cluster-types.md; delete this key if not needed>",
"NCCL_NET_PLUGIN": "none",
"NCCL_TIMEOUT": "900",
"OMP_NUM_THREADS": "8"
},
"debug": false,
"docker_gpus": "0,1,2,3,4,5,6,7"
}
}
124 changes: 124 additions & 0 deletions ...de/skills/mad-slurm-multinode/assets/manifests/primus_maxtext_llama3.1-405b.template.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,124 @@
{
"built_images": {
"rocm-primus-maxtext-llama3.1-405b-overlay": {
"model": "jax-maxtext/maxtext_MI355X_llama3.1_405B-fp8-pretrain",
"docker_image": "rocm/jax-training:maxtext-v26.6-<FILL_RCCL_TAG e.g. rccl-7.2.1 or rccl-develop-<sha7>>",
"dockerfile": "docker/primus_maxtext_rccl_overlay.ubuntu.amd.Dockerfile",
"base_docker": "rocm/jax-training:maxtext-v26.6",
"build_duration": 0,
"local_image": true,
"registry_image": null,
"registry": null,
"gpu_vendor": "AMD"
}
},
"built_models": {
"rocm-primus-maxtext-llama3.1-405b-overlay": {
"name": "jax-maxtext/maxtext_MI355X_llama3.1_405B-fp8-pretrain",
"url": "",
"dockerfile": "docker/primus_maxtext_rccl_overlay.ubuntu.amd.Dockerfile",
"scripts": "scripts/jax-maxtext/run.sh",
"n_gpus": "-1",
"owner": "mad.support@amd.com",
"training_precision": "fp8",
"multiple_results": "primus_perf_output.csv",
"tags": [
"maxtext",
"jax",
"MI355X",
"llama3.1_405B-fp8-pretrain",
"fp8",
"training"
],
"timeout": -1,
"args": "--config_path examples/maxtext/configs/MI355X/llama3.1_405B-fp8-pretrain.yaml per_device_batch_size=1",
"additional_docker_run_options": "--privileged --group-add render --shm-size 64G --device=/dev/infiniband --cap-add IPC_LOCK --ulimit memlock=-1 --ulimit nofile=1048576 -v /sys:/sys:ro -v /run/udev:/run/udev:ro"
}
},
"context": {
"docker_env_vars": {
"NCCL_DEBUG": "INFO",
Comment thread
i-kosarev marked this conversation as resolved.
Resolved
"NCCL_DEBUG_SUBSYS": "INIT,NET",
"NCCL_IB_DISABLE": "0",
"NCCL_NET": "IB",
"NCCL_IB_HCA": "<FILL_NCCL_IB_HCA data-plane HCA list, see references/cluster-types.md>",
"NCCL_IB_GID_INDEX": "<FILL_GID RoCEv2 GID index, see references/cluster-types.md>",
"NCCL_SOCKET_IFNAME": "<FILL_IFNAME management iface, see references/cluster-types.md>",
"GLOO_SOCKET_IFNAME": "<FILL_IFNAME management iface, see references/cluster-types.md>",
"RDMAV_DRIVERS": "<FILL_DRIVER see references/cluster-types.md>",
"IBV_DRIVERS": "<FILL_DRIVER see references/cluster-types.md>",
"RCCL_AINIC_ROCE": "<FILL_RCCL_AINIC_ROCE per archetype, see references/cluster-types.md; delete this key if not needed>",
"NCCL_NET_PLUGIN": "none",
"IBV_SHOW_WARNINGS": "1"
},
"docker_mounts": {
"/dev/infiniband": "/dev/infiniband"
},
"docker_build_arg": {
"BASE_DOCKER": "rocm/jax-training:maxtext-v26.6",
"RCCL_COMMIT": "<FILL_RCCL_COMMIT the commit to build; leave empty to track RCCL_BRANCH>",
"RCCL_BRANCH": "develop",
"RCCL_CMAKE_OPTIONS": "<FILL_RCCL_CMAKE_OPTIONS e.g. -DFAULT_INJECTION=ON; empty for a plain build>",
"BUILD_GPU_TARGETS": "<FILL_GPU_TARGETS gfx950 for MI350X/MI355X, gfx942 for MI300X/MI325X>",
"RDMA_CORE_VERSION": "<FILL_RDMA_CORE_VERSION 63.0 on Broadcom Thor2; leave empty to keep the base image's rdma-core>"
},
"gpu_vendor": "AMD",
"guest_os": "UBUNTU",
"docker_gpus": "0,1,2,3,4,5,6,7"
},
"credentials_required": [],
"summary": {
"successful_builds": [],
"failed_builds": [],
"total_build_time": 0,
"successful_pushes": [],
"failed_pushes": []
},
"deployment_config": {
"target": "slurm",
"slurm": {
"partition": "<FILL_PARTITION>",
"account": "<FILL_ACCOUNT or remove if cluster has none>",
"qos": "<FILL_QOS>",
"nodes": 4,
"gpus_per_node": 8,
"time": "12:00:00",
"output_dir": "./slurm_output",
"exclusive": true,
"network_interface": "<FILL_IFNAME management iface, see references/cluster-types.md>"
},
"distributed": {
"launcher": "primus",
"backend": "nccl",
"port": 29500,
"nnodes": 4,
"nproc_per_node": 8,
"primus": {
"config_path": "examples/maxtext/configs/MI355X/llama3.1_405B-fp8-pretrain.yaml",
"backend": "MaxText"
}
},
"env_vars": {
"HF_HOME": "${HF_HOME}",
"XDG_CACHE_HOME": "${XDG_CACHE_HOME}",
"PIP_CACHE_DIR": "${PIP_CACHE_DIR}",
"MAD_DATAHOME": "${MAD_DATAHOME}",
"NCCL_DEBUG": "INFO",
"NCCL_DEBUG_SUBSYS": "INIT,NET",
"NCCL_IB_DISABLE": "0",
"NCCL_NET": "IB",
"NCCL_SOCKET_IFNAME": "<FILL_IFNAME management iface, see references/cluster-types.md>",
"GLOO_SOCKET_IFNAME": "<FILL_IFNAME management iface, see references/cluster-types.md>",
"NCCL_IB_GID_INDEX": "<FILL_GID RoCEv2 GID index, see references/cluster-types.md>",
"NCCL_IB_HCA": "<FILL_NCCL_IB_HCA same value as context.docker_env_vars>",
"RDMAV_DRIVERS": "<FILL_DRIVER see references/cluster-types.md>",
"IBV_DRIVERS": "<FILL_DRIVER see references/cluster-types.md>",
"RCCL_AINIC_ROCE": "<FILL_RCCL_AINIC_ROCE per archetype, see references/cluster-types.md; delete this key if not needed>",
"NCCL_NET_PLUGIN": "none",
"NCCL_TIMEOUT": "900",
"OMP_NUM_THREADS": "8"
},
"debug": false,
"docker_gpus": "0,1,2,3,4,5,6,7"
}
}
Loading