Skip to content

Fix decoupled tests, refactor ICI parallelism override removal#3835

Open
gulsumgudukbay wants to merge 4 commits intoAI-Hypercomputer:mainfrom
ROCm:decoupled-tests-upstream
Open

Fix decoupled tests, refactor ICI parallelism override removal#3835
gulsumgudukbay wants to merge 4 commits intoAI-Hypercomputer:mainfrom
ROCm:decoupled-tests-upstream

Conversation

@gulsumgudukbay
Copy link
Copy Markdown
Collaborator

Description

Fix decoupled MaxText tests so they can run against the local decoupled test configuration without per-test ICI parallelism overrides.
Previously, many tests patched ici_fsdp_parallelism dynamically through a shared test helper. This made decoupled behavior harder to reason about and spread device-count-specific setup across unrelated tests. This PR moves that behavior into decoupled_base_test.yml by inheriting from base.yml and using config-level parallelism defaults (using @darisoy's changes in dbbc69e), then removes the ad-hoc helper from unit and integration tests.
The follow-up ROCm-specific fixes keep unsupported ROCm paths isolated: parameter host offload is skipped on ROCm, ROCm GPU architecture strings are not parsed as CUDA compute capabilities, and packed Transformer Engine attention passes the THD metadata required by newer TE versions while preserving compatibility with older versions.
This keeps the upstream-facing behavior focused on decoupled/ROCm test correctness rather than changing general test configs.

Tests

export PYTHONPATH=$PWD/src
export DECOUPLE_GCLOUD=TRUE
pytest -m 'not cpu_only and not tpu_only and not post_training and decoupled' -v tests --ignore=tests/post_training --csv=decoupled-tests-report.csv --html=decoupled-tests-report.html --self-contained-html 2>&1 | tee maxtext_decoupled_UT.log

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants