[DECOUPLED-MODE] Adding Decoupling Logic #2865
Conversation
(cherry picked from commit e8cc951)
(cherry picked from commit 0b58e96)
(cherry picked from commit 14f0508)
…ts library (cherry picked from commit 6f0b361)
(cherry picked from commit e43e370)
(cherry picked from commit 1c14d6c)
…ck, todo: remove this after updating jax. Configure ICI data parallelism for decoupled mode
gulsumgudukbay
requested review from
bvandermoon,
gobbleturk,
jacoguzo,
richjames0 and
shralex
as code owners
…re marks are used
__init__.py
Outdated
| @@ -0,0 +1,14 @@ | |||
| """Top-level shim for importing test_utils | |||
There was a problem hiding this comment.
License header please.
__init__.py
Outdated
| @@ -0,0 +1,14 @@ | |||
| """Top-level shim for importing test_utils | |||
| This shim lets test modules import `maxtext.tests`. | |||
There was a problem hiding this comment.
Can this __init__.py be moved inside /tests, since it contains the logic relevant for testing only?
|
|
||
| def test_tiny_config(self): | ||
| test_tmpdir = os.environ.get("TEST_TMPDIR") # pylint: disable=unused-variable | ||
| dataset_path = get_test_dataset_path() |
There was a problem hiding this comment.
This can be moved to a setUp() method, and it will be called automatically before each individual test methods.
| decoupled = is_decoupled() | ||
| dataset_path = get_test_dataset_path() | ||
| base_output_directory = ( | ||
| os.environ.get("LOCAL_BASE_OUTPUT", get_test_base_output_directory()) |
There was a problem hiding this comment.
Why don't we use this if logic in train_smoke_test.py?
| from MaxText.globals import MAXTEXT_PKG_DIR | ||
| from maxtext.tests.test_utils import get_test_config_path | ||
|
|
||
| pytestmark = [pytest.mark.tpu_only] |
There was a problem hiding this comment.
These tests are suppose to run on CPUs. Why are we adding tpu_only marker?
| from MaxText.input_pipeline import input_pipeline_interface | ||
| from maxtext.tests.test_utils import get_test_config_path, get_test_dataset_path, get_test_base_output_directory | ||
|
|
||
| MAXTEXT_ASSETS_ROOT = os.path.join("src", MAXTEXT_PKG_DIR, "assets") |
There was a problem hiding this comment.
MAXTEXT_ASSETS_ROOT is already initialized here: https://github.com/ROCm/maxtext/blob/4799cef7af51bb45272390e9b66bbae58790c0bd/src/MaxText/globals.py#L31
| _defects: list[str] = [] | ||
|
|
||
|
|
||
| def _import(name: str): |
There was a problem hiding this comment.
This method is not called anywhere, I guess?
2 participants
Description
This PR is the second part of the decoupling support. It adds logic for decoupling support, along with some test modifications for decoupling to be enabled.
Details:
Tests
All unit tests pass in decoupled mode.
UT results:
== 306 passed, 170 skipped, 25 deselected, 6588 warnings in 975.16s (0:16:15) ==
Train test:
python -m MaxText.train MaxText/configs/base.yml run_name=test hardware=gpu steps=5 model_name=llama2-7b attention=cudnn_flash_te enable_checkpointing=False ici_expert_parallelism=1 ici_fsdp_parallelism=-1 ici_data_parallelism=1 remat_policy=minimal scan_layers=True dataset_type=synthetic logits_dot_in_fp32=False dtype=bfloat16 weight_dtype=bfloat16 per_device_batch_size=1 max_target_length=2048 shardy=False
works.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.