File size: 1,181 Bytes
dfb775d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | """FSDP / DeepSpeed config builders."""
from __future__ import annotations
import pytest
from mindxtrain.train.distributed import build_deepspeed_config, build_fsdp_config
def test_fsdp_8gpu_config_has_correct_strategy():
cfg = build_fsdp_config(8)
assert cfg["num_processes"] == 8
assert cfg["distributed_type"] == "FSDP"
assert cfg["fsdp_config"]["fsdp_sharding_strategy"] == "FULL_SHARD"
def test_fsdp_1gpu_config_is_no_distributed():
cfg = build_fsdp_config(1)
assert cfg["num_processes"] == 1
assert cfg["distributed_type"] == "NO"
def test_fsdp_2gpu_rejected_xgmi():
with pytest.raises(ValueError, match="1 or 8"):
build_fsdp_config(2) # type: ignore[arg-type]
def test_fsdp_4gpu_rejected_xgmi():
with pytest.raises(ValueError, match="1 or 8"):
build_fsdp_config(4) # type: ignore[arg-type]
def test_deepspeed_zero3_default():
cfg = build_deepspeed_config()
assert cfg["zero_optimization"]["stage"] == 3
assert cfg["bf16"]["enabled"] is True
def test_deepspeed_invalid_stage_rejected():
with pytest.raises(ValueError):
build_deepspeed_config(zero_stage=4) # type: ignore[arg-type]
|