[DECOUPLED-MODE] Adding Decoupling Logic#2865

Open

gulsumgudukbay wants to merge 26 commits intoAI-Hypercomputer:mainfrom

ROCm:rocm-main-pr

Commits on Dec 21, 2025

adding necessary files
gulsumgudukbay
committed
add decoupling logic config patch to tests
gulsumgudukbay
committed
add correct ICI parallelism to tests for decoupled mode
gulsumgudukbay
committed
Add tpu_only marker to train compile test
gulsumgudukbay
committed
fixing more UTs
gulsumgudukbay
committed
adding decoupling logic, biggest change
gulsumgudukbay
committed
add tensorboardX stub
gulsumgudukbay
committed
adding tokamax changes along with some UT fixes
gulsumgudukbay
committed
fixing little UT issues:
gulsumgudukbay
committed
fixing train_tests
gulsumgudukbay
committed
removing CI workflows for now to upstream decoupling changes
gulsumgudukbay
committed
making jetstream and tunix optional and add is_stub variables
gulsumgudukbay
committed
removing tunix from decoupling logic
gulsumgudukbay
committed
removing tunix from decoupled mode logic
gulsumgudukbay
committed
addressing PR comments
gulsumgudukbay
committed
fixing pylint issues
gulsumgudukbay
committed
renaming datasets to local_datasets to avoid confusion with HF datasets library
gulsumgudukbay
committed
pyink fixes
gulsumgudukbay
committed
Rename GCE_MARKERS to GCP_MARKERS
gulsumgudukbay
committed
updating dataset paths and fix gcloud_Stub
gulsumgudukbay
committed
making decoupled mode work with upstream updates
gulsumgudukbay
committed
updates for upstream sync UTs
gulsumgudukbay
committed
remove context_parallel_strategy config param, todo: add it back later
gulsumgudukbay
committed
make jax_remove_size_one_mesh_axis_from_type param setting in try block, todo: remove this after updating jax. Configure ICI data parallelism for decoupled mode
gulsumgudukbay
committed
Fix decoupled rampup and mesh configs for bare-metal tests
gulsumgudukbay
committed
remove some stuff from this PR that were meant to be for downstream
gulsumgudukbay
committed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[DECOUPLED-MODE] Adding Decoupling Logic#2865

[DECOUPLED-MODE] Adding Decoupling Logic#2865
gulsumgudukbay wants to merge 26 commits intoAI-Hypercomputer:mainfrom
ROCm:rocm-main-pr

Commits on Dec 21, 2025

adding necessary files

add decoupling logic config patch to tests

add correct ICI parallelism to tests for decoupled mode

Add tpu_only marker to train compile test

fixing more UTs

adding decoupling logic, biggest change

add tensorboardX stub

adding tokamax changes along with some UT fixes

fixing little UT issues:

fixing train_tests

removing CI workflows for now to upstream decoupling changes

making jetstream and tunix optional and add is_stub variables

removing tunix from decoupling logic

removing tunix from decoupled mode logic

addressing PR comments

fixing pylint issues

renaming datasets to local_datasets to avoid confusion with HF datasets library

pyink fixes

Rename GCE_MARKERS to GCP_MARKERS

updating dataset paths and fix gcloud_Stub

making decoupled mode work with upstream updates

updates for upstream sync UTs

remove context_parallel_strategy config param, todo: add it back later

make jax_remove_size_one_mesh_axis_from_type param setting in try block, todo: remove this after updating jax. Configure ICI data parallelism for decoupled mode

Fix decoupled rampup and mesh configs for bare-metal tests

remove some stuff from this PR that were meant to be for downstream

[DECOUPLED-MODE] Adding Decoupling Logic#2865gulsumgudukbay wants to merge 26 commits intoAI-Hypercomputer:mainAI-Hypercomputer/maxtext:mainfrom ROCm:rocm-main-prROCm/maxtext:rocm-main-prCopy head branch name to clipboard

Commits

Commits on Dec 21, 2025

[DECOUPLED-MODE] Adding Decoupling Logic#2865
gulsumgudukbay wants to merge 26 commits intoAI-Hypercomputer:mainfrom
ROCm:rocm-main-pr