diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py new file mode 100644 index 0000000000..bf9db094d2 --- /dev/null +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -0,0 +1,85 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared pytest fixtures for FSDP2 distributed tests. + +Fixtures defined here (dist_init, _cleanup, recipe_name) are auto-discovered +by pytest for every test module in this directory. +""" + +import gc +import os +import pytest +import torch +import torch.distributed as dist +from transformer_engine.pytorch import fp8 + +# Ensure the correct CUDA device is active before _parametrize_recipes() +# runs at collection time, since the session-scoped dist_init fixture +# has not executed yet. +_local_rank = int(os.environ.get("LOCAL_RANK", "0")) +torch.cuda.set_device(_local_rank) + + +# ── FP8 recipe parametrization ────────────────────────────────────── +def _check_nvfp4_support(): + supported, reason = fp8.check_nvfp4_support() + if supported and torch.cuda.get_device_capability()[0] == 12: + return ( + False, + ( + "NVFP4BlockScaling is failing on SM120 with " + "hadamard_transform/hadamard_transform_cast_fusion.cu:672 in function " + "rht_gemm_ntt_w_sfc: CUDA Error: invalid argument" + ), + ) + return supported, reason + + +_FP8_RECIPE_CONFIGS = [ + ("DelayedScaling", fp8.check_fp8_support), + ("Float8CurrentScaling", fp8.check_fp8_support), + ("Float8BlockScaling", fp8.check_fp8_block_scaling_support), + ("MXFP8BlockScaling", fp8.check_mxfp8_support), + ("NVFP4BlockScaling", _check_nvfp4_support), +] + + +def _parametrize_recipes(): + params = [] + for name, check_fn in _FP8_RECIPE_CONFIGS: + supported, reason = check_fn() + params.append( + pytest.param(name, id=name, marks=pytest.mark.skipif(not supported, reason=reason)) + ) + return params + + +# ── Session / per-test fixtures ────────────────────────────────────── +@pytest.fixture(scope="session", autouse=True) +def dist_init(): + """Initialize the distributed process group once for the entire pytest session.""" + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="cpu:gloo,cuda:nccl") + torch.manual_seed(42) + torch.cuda.manual_seed(42) + yield + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.fixture(autouse=True) +def _cleanup(): + """Release GPU memory and stale NCCL state between tests.""" + yield + if dist.is_initialized(): + dist.barrier() + gc.collect() + torch.cuda.empty_cache() + + +@pytest.fixture(params=_parametrize_recipes()) +def recipe_name(request): + return request.param diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py new file mode 100644 index 0000000000..178ce62375 --- /dev/null +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -0,0 +1,31 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared utility functions for FSDP2 distributed tests.""" + +import transformer_engine.common.recipe +from transformer_engine.pytorch import QuantizedTensor + + +def get_recipe_from_string(recipe): + return getattr(transformer_engine.common.recipe, recipe)() + + +def save_custom_attrs(module): + custom_attrs = {} + for name, param in module.named_parameters(): + if isinstance(param, QuantizedTensor): + ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] + else: + ignore_keys = [] + attrs = vars(param) + custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} + return custom_attrs + + +def restore_custom_attrs(module, custom_attrs): + for name, param in module.named_parameters(): + if name in custom_attrs: + for attr_name, attr_value in custom_attrs[name].items(): + setattr(param, attr_name, attr_value) diff --git a/tests/pytorch/distributed/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py similarity index 58% rename from tests/pytorch/distributed/run_fsdp2_fused_adam.py rename to tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index c39957cf13..877fa66795 100644 --- a/tests/pytorch/distributed/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -6,12 +6,28 @@ """FSDP2 + FusedAdam compatibility tests. -Launched via torchrun from test_fused_optimizer.py. +Run all tests (via torchrun + pytest): + torchrun -m pytest -v --tb=short + +Run a single test standalone (for debugging): + torchrun --test --recipe + +Available --test values: + fused_adam_fp8_master_weights, fused_adam_fp8_master_weights_no_meta, + fused_adam_bf16, fused_adam_fp8_no_master, fused_adam_bf16_store_param_remainders, + fuse_wgrad_accumulation, dcp_output_parity, dcp_output_parity_async, + safetensors_fp32_export + +Available --recipe values: + DelayedScaling, Float8CurrentScaling, Float8BlockScaling, + MXFP8BlockScaling, NVFP4BlockScaling """ import argparse import functools import os +import shutil +import pytest import torch import torch.distributed as dist @@ -24,9 +40,7 @@ from transformer_engine.pytorch import QuantizedTensor import transformer_engine.common.recipe - -def get_recipe_from_string(recipe): - return getattr(transformer_engine.common.recipe, recipe)() +from fsdp2_utils import get_recipe_from_string, save_custom_attrs, restore_custom_attrs HIDDEN_SIZE = 256 @@ -38,38 +52,6 @@ def get_recipe_from_string(recipe): NUM_STEPS = 3 -def save_custom_attrs(module): - custom_attrs = {} - for name, param in module.named_parameters(): - if isinstance(param, QuantizedTensor): - ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] - else: - ignore_keys = [] - attrs = vars(param) - custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} - return custom_attrs - - -def restore_custom_attrs(module, custom_attrs): - for name, param in module.named_parameters(): - if name in custom_attrs: - for attr_name, attr_value in custom_attrs[name].items(): - setattr(param, attr_name, attr_value) - - -def _setup(): - """Common distributed setup. Returns (world_size, local_rank, device).""" - world_size = int(os.environ["WORLD_SIZE"]) - local_rank = int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) - # CPU backend required for async save - dist.init_process_group(backend="cpu:gloo,cuda:nccl") - device = torch.device(f"cuda:{local_rank}") - torch.manual_seed(42) - torch.cuda.manual_seed(42) - return world_size, local_rank, device - - def _build_model(fp8_init, fuse_wgrad_accumulation=False, recipe=None, use_meta_device=True): """Build a Sequential of TransformerLayers, optionally with FP8 init. @@ -143,7 +125,14 @@ def _shard_model(model, world_size): return model -def test_fused_adam_fp8_master_weights(recipe=None): +def _get_dist_info(): + """Get world_size and device from environment (PG already initialized by session fixture).""" + world_size = int(os.environ["WORLD_SIZE"]) + device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}") + return world_size, device + + +def test_fused_adam_fp8_master_weights(recipe_name): """FusedAdam with master_weights + FSDP2 + quantized_model_init (FP8 params). Verifies: @@ -151,7 +140,15 @@ def test_fused_adam_fp8_master_weights(recipe=None): - Training loop completes without error - DTensor wrapping and QuantizedTensor local tensors are preserved """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + + if recipe_name == "NVFP4BlockScaling": + pytest.xfail( + f"{recipe_name}: quantized_model_init and FSDP2 is not currently supported, since the " + "block tensor is dequantized before we flatten it for FSDP2." + ) + + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe) model = _shard_model(model, world_size) @@ -206,10 +203,8 @@ def test_fused_adam_fp8_master_weights(recipe=None): ) assert qt_count > 0, "No QuantizedTensor local tensors after training" - dist.destroy_process_group() - -def test_fused_adam_fp8_master_weights_no_meta(recipe=None): +def test_fused_adam_fp8_master_weights_no_meta(recipe_name): """FusedAdam with master_weights + FSDP2 + quantized_model_init WITHOUT meta device. This is the legacy path that creates quantized params directly on CUDA. @@ -219,7 +214,16 @@ def test_fused_adam_fp8_master_weights_no_meta(recipe=None): For per-tensor FP8 (DelayedScaling, Float8CurrentScaling) this works because Float8Tensor's storage is accessible via data_ptr(). """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + + if recipe_name in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): + pytest.xfail( + f"{recipe_name}: FSDP2 without meta-device init crashes on block-scaling " + "QuantizedTensor wrapper subclasses (data_ptr() == 0). " + "Use device='meta' + reset_parameters() after sharding." + ) + + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe, use_meta_device=False) model = _shard_model(model, world_size) @@ -242,15 +246,15 @@ def test_fused_adam_fp8_master_weights_no_meta(recipe=None): loss.backward() optimizer.step() - dist.destroy_process_group() - -def test_fused_adam_bf16(recipe=None): +def test_fused_adam_bf16(recipe_name): """FusedAdam with master_weights + FSDP2 + bf16 params (no FP8). Verifies the non-FP8 DTensor param path in step() works correctly. """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + + world_size, device = _get_dist_info() model = _build_model(fp8_init=False) model = _shard_model(model, world_size) @@ -284,15 +288,21 @@ def test_fused_adam_bf16(recipe=None): # Verify loss decreased (basic sanity) assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" - dist.destroy_process_group() - -def test_fused_adam_fp8_no_master(recipe=None): +def test_fused_adam_fp8_no_master(recipe_name): """FusedAdam without master_weights + FSDP2 + FP8 params. Verifies FusedAdam works with FSDP2 even without master weights enabled. """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + + if recipe_name in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): + pytest.xfail( + f"{recipe_name}: FusedAdam without master_weights does not support " + "block-scaling quantized tensors. Use master_weights=True." + ) + + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe) model = _shard_model(model, world_size) @@ -318,10 +328,8 @@ def test_fused_adam_fp8_no_master(recipe=None): for name, param in model.named_parameters(): assert isinstance(param, DTensor), f"{name} lost DTensor wrapping" - dist.destroy_process_group() - -def test_fused_adam_bf16_store_param_remainders(recipe=None): +def test_fused_adam_bf16_store_param_remainders(recipe_name): """FusedAdam with master_weights + store_param_remainders + FSDP2 + bf16 params. store_param_remainders stores only the trailing 16 remainder bits (int16) @@ -335,7 +343,8 @@ def test_fused_adam_bf16_store_param_remainders(recipe=None): - exp_avg and exp_avg_sq are float32 - Loss decreases (basic sanity) """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + world_size, device = _get_dist_info() model = _build_model(fp8_init=False) model = _shard_model(model, world_size) @@ -385,10 +394,18 @@ def test_fused_adam_bf16_store_param_remainders(recipe=None): # Verify loss decreased (basic sanity) assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" - dist.destroy_process_group() - -def test_fuse_wgrad_accumulation(recipe=None): +@pytest.mark.xfail( + reason=( + "fuse_wgrad_accumulation is incompatible with vanilla FSDP2: " + "autograd Function.apply unwraps DTensors to local tensors, so " + "main_grad (set on the DTensor) is inaccessible during backward. " + "Additionally, the fused wgrad GEMM bypasses FSDP2's reduce-scatter." + ), + raises=AttributeError, + strict=True, +) +def test_fuse_wgrad_accumulation(recipe_name): """fuse_wgrad_accumulation=True + FSDP2 -- expected to fail. With vanilla FSDP2, PyTorch's autograd Function.apply unwraps DTensor @@ -400,8 +417,8 @@ def test_fuse_wgrad_accumulation(recipe=None): writes the gradient directly into main_grad and returns None to autograd, bypassing FSDP2's reduce-scatter. """ - world_size, _, device = _setup() - + recipe = get_recipe_from_string(recipe_name) + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, fuse_wgrad_accumulation=True, recipe=recipe) # Allocate main_grad buffers on the DTensor params @@ -433,10 +450,8 @@ def test_fuse_wgrad_accumulation(recipe=None): loss = F.mse_loss(output, target) loss.backward() # Expected to raise AttributeError - dist.destroy_process_group() - -def test_safetensors_fp32_export(recipe=None): +def test_safetensors_fp32_export(recipe_name): """Export full-precision (FP32) model to safetensors from optimizer master weights. Verifies: @@ -446,6 +461,13 @@ def test_safetensors_fp32_export(recipe=None): - All saved tensors are float32 - Saved tensor shapes match expected (unsharded) shapes """ + recipe = get_recipe_from_string(recipe_name) + if recipe_name == "MXFP8BlockScaling": + pytest.xfail( + "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " + "MXFP8 quantized tensors, causing illegal memory access" + ) + from safetensors.torch import load_file, save_file from torch.distributed.checkpoint.state_dict import ( StateDictOptions, @@ -453,8 +475,7 @@ def test_safetensors_fp32_export(recipe=None): get_optimizer_state_dict, ) - world_size, _, device = _setup() - + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe) model = _shard_model(model, world_size) @@ -483,38 +504,39 @@ def test_safetensors_fp32_export(recipe=None): full_opt_state = get_optimizer_state_dict(model, optimizer, options=full_opts) rank = int(os.environ.get("RANK", "0")) - save_path = "/tmp/te_test_fsdp2_model_fp32.safetensors" + save_path = f"/tmp/te_test_fsdp2_model_fp32_{recipe_name}.safetensors" if rank == 0: - # Build FP32 state dict from optimizer master weights. - fp32_state = {} - opt_param_states = full_opt_state.get("state", {}) - - for key, value in full_model_state.items(): - if key in opt_param_states and "master_param" in opt_param_states[key]: - fp32_state[key] = opt_param_states[key]["master_param"].float() - else: - fp32_state[key] = value.float() + if os.path.exists(save_path): + os.remove(save_path) - assert len(fp32_state) > 0, "FP32 state dict is empty" + try: + fp32_state = {} + opt_param_states = full_opt_state.get("state", {}) - # Save and verify. - save_file(fp32_state, save_path) - loaded = load_file(save_path) + for key, value in full_model_state.items(): + if key in opt_param_states and "master_param" in opt_param_states[key]: + fp32_state[key] = opt_param_states[key]["master_param"].float() + else: + fp32_state[key] = value.float() - assert len(loaded) == len( - fp32_state - ), f"Loaded {len(loaded)} tensors, expected {len(fp32_state)}" - for k, v in loaded.items(): - assert v.dtype == torch.float32, f"{k}: expected float32, got {v.dtype}" + assert len(fp32_state) > 0, "FP32 state dict is empty" - # Clean up. - os.remove(save_path) + save_file(fp32_state, save_path) + loaded = load_file(save_path) - dist.destroy_process_group() + assert len(loaded) == len( + fp32_state + ), f"Loaded {len(loaded)} tensors, expected {len(fp32_state)}" + for k, v in loaded.items(): + assert v.dtype == torch.float32, f"{k}: expected float32, got {v.dtype}" + finally: + if os.path.exists(save_path): + os.remove(save_path) -def test_dcp_output_parity(recipe=None, async_save=False): +@pytest.mark.parametrize("async_save", [False, True], ids=["sync", "async"]) +def test_dcp_output_parity(recipe_name, async_save): """DCP save/load round-trip produces bitwise-identical model outputs. 1. Builds and trains a model for NUM_STEPS @@ -525,156 +547,197 @@ def test_dcp_output_parity(recipe=None, async_save=False): 6. Runs the same forward pass and asserts outputs are identical 7. Runs one more training step on both models and asserts outputs still match """ - import torch.distributed.checkpoint as dcp - - world_size, local_rank, device = _setup() - - # ── Build and train the original model ─────────────────────────── - model = _build_model(fp8_init=True, recipe=recipe) - model = _shard_model(model, world_size) - - optimizer = te.optimizers.FusedAdam( - model.parameters(), - lr=1e-3, - master_weights=True, - master_weight_dtype=torch.float32, - ) - - x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) - target = torch.randn_like(x) - - for _ in range(NUM_STEPS): - optimizer.zero_grad(set_to_none=True) - with te.autocast(enabled=True, recipe=recipe): - output = model(x) - loss = F.mse_loss(output, target) - loss.backward() - optimizer.step() + recipe = get_recipe_from_string(recipe_name) + + if recipe_name == "MXFP8BlockScaling": + pytest.xfail( + "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " + "MXFP8 quantized tensors, causing illegal memory access: " + "/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh:92 in function " + "multi_tensor_apply: CUDA Error: an illegal memory access was encountered" + ) - # Record reference output from the trained model. - with torch.no_grad(): - with te.autocast(enabled=True, recipe=recipe): - ref_output = model(x).clone() - - # ── Save checkpoint ────────────────────────────────────────────── - checkpoint_dir = "/tmp/te_test_fsdp2_dcp_parity" - - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): - # We need to remove the _extra_state keys from the model state dict for DelayedScaling, - # since otherwise we'll run into an error that the tensor sizes are different. The - # alternative is a LoadPlanner that dynamically re-sizes the input tensors, see - # NVIDIA/TransformerEngine#1860 for more details. - model_state = { - k: v for k, v in model.state_dict().items() if not k.endswith("_extra_state") - } - else: - model_state = model.state_dict() + if recipe_name == "NVFP4BlockScaling": + pytest.xfail( + "NVFP4BlockScaling: DCP load_state_dict triggers reset_sharded_param() " + "which calls data_ptr() on NVFP4Tensor wrapper subclass with invalid storage" + ) - save_state = {"model": model_state, "optimizer": optimizer.state_dict()} + if ( + recipe_name == "Float8BlockScaling" + and not async_save + and torch.cuda.get_device_capability()[0] == 12 + ): + pytest.xfail( + "Float8BlockScaling is failing on SM120 with RuntimeError: " + "transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu:534 " + "in function quantize_transpose_vector_blockwise: Assertion failed: pow2_scale. On " + "Blackwell and newer, the FP8 block scaling recipe is emulated with MXFP8, which " + "requires using power of two scaling factors." + ) + if recipe_name == "Float8BlockScaling" and async_save: + pytest.xfail( + "Float8BlockScaling: async DCP save/load round-trip produces different model " + "outputs — quantization metadata (scales) is not correctly persisted through " + "async distributed checkpointing. On SM120, additionally fails with pow2_scale " + "assertion in quantize_transpose_vector_blockwise." + ) - if not async_save: - dcp.save(save_state, checkpoint_id=checkpoint_dir) - else: - future = dcp.async_save(save_state, checkpoint_id=checkpoint_dir) - future.result() # Block on async save completion + import torch.distributed.checkpoint as dcp - # ── Build a fresh model and load the checkpoint ────────────────── - model2 = _build_model(fp8_init=True, recipe=recipe) - model2 = _shard_model(model2, world_size) + world_size, device = _get_dist_info() + rank = int(os.environ.get("RANK", "0")) + save_mode = "async" if async_save else "sync" + checkpoint_dir = f"/tmp/te_test_fsdp2_dcp_parity_{recipe_name}_{save_mode}" - optimizer2 = te.optimizers.FusedAdam( - model2.parameters(), - lr=1e-3, - master_weights=True, - master_weight_dtype=torch.float32, - ) + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + dist.barrier() + + try: + # ── Build and train the original model ─────────────────────────── + model = _build_model(fp8_init=True, recipe=recipe) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) - # Populate optimizer state so load_state_dict has matching structure. - optimizer2.zero_grad(set_to_none=True) - with te.autocast(enabled=True, recipe=recipe): - out_tmp = model2(x) - F.mse_loss(out_tmp, target).backward() - optimizer2.step() - - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): - model2_state = { - k: v for k, v in model2.state_dict().items() if not k.endswith("_extra_state") - } - else: - model2_state = model2.state_dict() + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + + # Record reference output from the trained model. + with torch.no_grad(): + with te.autocast(enabled=True, recipe=recipe): + ref_output = model(x).clone() + + # ── Save checkpoint ────────────────────────────────────────────── + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + # We need to remove the _extra_state keys from the model state dict for + # DelayedScaling, since otherwise we'll run into an error that the tensor + # sizes are different. The alternative is a LoadPlanner that dynamically + # re-sizes the input tensors, see NVIDIA/TransformerEngine#1860 for more + # details. + model_state = { + k: v for k, v in model.state_dict().items() if not k.endswith("_extra_state") + } + else: + model_state = model.state_dict() - state_to_load = {"model": model2_state, "optimizer": optimizer2.state_dict()} + save_state = {"model": model_state, "optimizer": optimizer.state_dict()} - dcp.load(state_to_load, checkpoint_id=checkpoint_dir) - model2.load_state_dict( - state_to_load["model"], - strict=( - False if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling) else True - ), - ) - optimizer2.load_state_dict(state_to_load["optimizer"]) - - # ── Verify identical forward-pass output ───────────────────────── - with torch.no_grad(): - with te.autocast(enabled=True, recipe=recipe): - loaded_output = model2(x) - - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): - # DelayedScaling stores amax history and scaling factors in _extra_state, - # which cannot be saved via DCP due to non-deterministic pickle sizes - # across ranks. The fresh model therefore uses default scaling factors, - # producing small numerical differences from FP8 re-quantization. - torch.testing.assert_close( - loaded_output, - ref_output, - rtol=0.05, - atol=0.1, - msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", - ) - else: - torch.testing.assert_close( - loaded_output, - ref_output, - rtol=0, - atol=0, - msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", + if not async_save: + dcp.save(save_state, checkpoint_id=checkpoint_dir) + else: + future = dcp.async_save(save_state, checkpoint_id=checkpoint_dir) + future.result() + + # ── Build a fresh model and load the checkpoint ────────────────── + model2 = _build_model(fp8_init=True, recipe=recipe) + model2 = _shard_model(model2, world_size) + + optimizer2 = te.optimizers.FusedAdam( + model2.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, ) - # ── Verify one more training step produces identical results ───── - optimizer.zero_grad(set_to_none=True) - with te.autocast(enabled=True, recipe=recipe): - out1 = model(x) - loss1 = F.mse_loss(out1, target) - loss1.backward() - optimizer.step() - - optimizer2.zero_grad(set_to_none=True) - with te.autocast(enabled=True, recipe=recipe): - out2 = model2(x) - loss2 = F.mse_loss(out2, target) - loss2.backward() - optimizer2.step() - - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): - torch.testing.assert_close( - out2, - out1, - rtol=0.05, - atol=0.1, - msg="Training step after DCP load produces different output", - ) - else: - torch.testing.assert_close( - out2, out1, msg="Training step after DCP load produces different output" + # Populate optimizer state so load_state_dict has matching structure. + optimizer2.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + out_tmp = model2(x) + F.mse_loss(out_tmp, target).backward() + optimizer2.step() + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + model2_state = { + k: v for k, v in model2.state_dict().items() if not k.endswith("_extra_state") + } + else: + model2_state = model2.state_dict() + + state_to_load = {"model": model2_state, "optimizer": optimizer2.state_dict()} + + dcp.load(state_to_load, checkpoint_id=checkpoint_dir) + model2.load_state_dict( + state_to_load["model"], + strict=( + False + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling) + else True + ), ) + optimizer2.load_state_dict(state_to_load["optimizer"]) + + # ── Verify identical forward-pass output ───────────────────────── + with torch.no_grad(): + with te.autocast(enabled=True, recipe=recipe): + loaded_output = model2(x) + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + # DelayedScaling stores amax history and scaling factors in _extra_state, + # which cannot be saved via DCP due to non-deterministic pickle sizes + # across ranks. The fresh model therefore uses default scaling factors, + # producing small numerical differences from FP8 re-quantization. + torch.testing.assert_close( + loaded_output, + ref_output, + rtol=0.05, + atol=0.1, + msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", + ) + else: + torch.testing.assert_close( + loaded_output, + ref_output, + rtol=0, + atol=0, + msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", + ) + + # ── Verify one more training step produces identical results ───── + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + out1 = model(x) + loss1 = F.mse_loss(out1, target) + loss1.backward() + optimizer.step() - # ── Cleanup ────────────────────────────────────────────────────── - import shutil - - if int(os.environ.get("RANK", "0")) == 0: - shutil.rmtree(checkpoint_dir, ignore_errors=True) - - dist.destroy_process_group() + optimizer2.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + out2 = model2(x) + loss2 = F.mse_loss(out2, target) + loss2.backward() + optimizer2.step() + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + torch.testing.assert_close( + out2, + out1, + rtol=0.05, + atol=0.1, + msg="Training step after DCP load produces different output", + ) + else: + torch.testing.assert_close( + out2, out1, msg="Training step after DCP load produces different output" + ) + finally: + dist.barrier() + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) TESTS = { @@ -707,5 +770,13 @@ def test_dcp_output_parity(recipe=None, async_save=False): ], ) args = parser.parse_args() - recipe = get_recipe_from_string(args.recipe) - TESTS[args.test](recipe) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="cpu:gloo,cuda:nccl") + torch.manual_seed(42) + torch.cuda.manual_seed(42) + try: + TESTS[args.test](args.recipe) + finally: + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/tests/pytorch/distributed/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py similarity index 80% rename from tests/pytorch/distributed/run_fsdp2_model.py rename to tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 60d7cd2023..fce565ed9a 100644 --- a/tests/pytorch/distributed/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -4,9 +4,36 @@ # # See LICENSE for license information. +"""FSDP2 model sharding tests. + +Run all tests (via torchrun + pytest): + torchrun -m pytest -v --tb=short + +Run standalone (for debugging): + torchrun --recipe [options] + +Available --recipe values: + DelayedScaling, Float8CurrentScaling, Float8BlockScaling, + MXFP8BlockScaling, NVFP4BlockScaling + +Other options: + --fp8-init Initialize weights in FP8 + --layer-type TYPE Linear, LayerNormLinear, LayerNormMLP, + MultiheadAttention, TransformerLayer (default) + --sharding-dims N [M] FSDP dims, e.g. "2" or "2 2" for HSDP + --num-layers N Number of layers (default: 4) + --iter N Training iterations (default: 10) + --device cuda|meta Device for init (default: meta) +""" + +import gc import os import sys import argparse +from types import SimpleNamespace +from contextlib import nullcontext + +import pytest import transformer_engine.pytorch as te import transformer_engine.common.recipe @@ -19,14 +46,12 @@ from torch.distributed import DeviceMesh from torch.distributed._composable.fsdp import fully_shard from torch.distributed.device_mesh import init_device_mesh -from transformer_engine.pytorch import QuantizedTensor -from contextlib import nullcontext -LOCAL_RANK = None +from fsdp2_utils import get_recipe_from_string, save_custom_attrs, restore_custom_attrs def dist_print(msg): - if LOCAL_RANK == 0: + if int(os.getenv("LOCAL_RANK", "0")) == 0: print(msg) @@ -114,10 +139,6 @@ def get_te_layer_from_string(layer_name): return te_layer_map[layer_name.lower()] -def get_recipe_from_string(recipe): - return getattr(transformer_engine.common.recipe, recipe)() - - def init_te_model(config): hidden_size = config.num_heads * config.head_dim args = [hidden_size, hidden_size] @@ -188,31 +209,8 @@ def shard_model_with_fsdp2(model, mesh): return model -#### Methods to save the custom attributes of QuantizedTensors before sharding -#### them with FSDP2, and restore them after sharding. -def save_custom_attrs(module): - custom_attrs = {} - for name, param in module.named_parameters(): - if isinstance(param, QuantizedTensor): - # Ignore FP8 metadata attributes. Otherwise we will save duplicate copies - # for data/transpose FP8 tensors on top of FP8 tensors that FSDP2 will save. - ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] - else: - ignore_keys = [] - attrs = vars(param) - custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} - return custom_attrs - - -def restore_custom_attrs(module, custom_attrs): - for name, param in module.named_parameters(): - if name in custom_attrs: - for attr_name, attr_value in custom_attrs[name].items(): - setattr(param, attr_name, attr_value) - - @torch.no_grad() -def test_fp8_fsdp2_allgather(model): +def _check_fp8_fsdp2_allgather(model): # Do manual allgather in fp32 and match against fp8 allgather done # with fsdp2 # FP32 manual weight allgather @@ -249,30 +247,10 @@ def test_fp8_fsdp2_allgather(model): module.reshard() -def _train(args): - global LOCAL_RANK - assert "TORCHELASTIC_RUN_ID" in os.environ - WORLD_RANK = int(os.getenv("RANK", "0")) - WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) - LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) - LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", "1")) - assert LOCAL_SIZE == WORLD_SIZE - - # Set device and initialize RNG states - torch.cuda.set_device(WORLD_RANK) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - # Initialize torch.distributed global process group and get DP/TP groups - dist_init_kwargs = { - "backend": "nccl", - "rank": WORLD_RANK, - "world_size": WORLD_SIZE, - } - assert dist.is_nccl_available() - dist.init_process_group(**dist_init_kwargs) - nccl_world = dist.new_group(backend="nccl") - device = torch.device(f"cuda:{LOCAL_RANK}") +def _run_training(args): + """Core training logic. Assumes dist is already initialized.""" + device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") + world_size = int(os.getenv("WORLD_SIZE", "1")) # FP8 Configuration fp8_recipe = get_recipe_from_string(args.recipe) @@ -298,7 +276,6 @@ def _train(args): ) # Creating a DeviceMesh for fully_shard - world_size = int(WORLD_SIZE) # Setup the sharding mesh for FSDP/HSDP mesh = get_device_mesh(world_size, args.sharding_dims) custom_attrs = save_custom_attrs(model) @@ -344,11 +321,71 @@ def _train(args): # Some of the FSDP states are lazy initialized during FSDP forward pass # so testing fp8 allgather at the end of the training loop. if args.fp8_init: - test_fp8_fsdp2_allgather(model) + _check_fp8_fsdp2_allgather(model) + + +def _train(args): + """Standalone entry point with full dist lifecycle.""" + assert "TORCHELASTIC_RUN_ID" in os.environ + WORLD_RANK = int(os.getenv("RANK", "0")) + WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) + LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) + LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", "1")) + assert LOCAL_SIZE == WORLD_SIZE + + torch.cuda.set_device(LOCAL_RANK) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + assert dist.is_nccl_available() + dist.init_process_group( + backend="nccl", + rank=WORLD_RANK, + world_size=WORLD_SIZE, + ) + try: + _run_training(args) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + torch.cuda.empty_cache() + gc.collect() - dist.destroy_process_group() return 0 +# ── Pytest test function ───────────────────────────────────────────── + +NUM_PROCS = int(os.environ.get("WORLD_SIZE", "1")) + + +@pytest.mark.parametrize("sharding_dims", [[NUM_PROCS], [2, NUM_PROCS // 2]]) +@pytest.mark.parametrize("fp8_init", [False, True]) +@pytest.mark.parametrize("layer_type", ["LayerNormLinear", "TransformerLayer"]) +def test_distributed(recipe_name, fp8_init, sharding_dims, layer_type): + if recipe_name in ("Float8BlockScaling", "NVFP4BlockScaling") and fp8_init: + pytest.xfail(f"{recipe_name} + fp8_init: test_fp8_fsdp2_allgather is currently failing.") + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + args = SimpleNamespace( + recipe=recipe_name, + fp8_init=fp8_init, + sharding_dims=list(sharding_dims), + layer_type=layer_type, + seed=42, + num_heads=8, + head_dim=64, + batch_size=16, + seq_length=128, + params_dtype="float32", + num_layers=4, + iter=10, + device="meta", + ) + _run_training(args) + + if __name__ == "__main__": sys.exit(_train(_parse_args())) diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index 02e45d99cb..aca8d6d692 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -10,242 +10,56 @@ import torch import transformer_engine.pytorch as te -from transformer_engine.pytorch import fp8 NUM_PROCS: int = torch.cuda.device_count() - - -def check_nvfp4_support(): - supported, reason = fp8.check_nvfp4_support() - if supported and torch.cuda.get_device_capability()[0] == 12: - return ( - False, - ( - "NVFP4BlockScaling is failing on SM120 with " - "hadamard_transform/hadamard_transform_cast_fusion.cu:672 in function " - "rht_gemm_ntt_w_sfc: CUDA Error: invalid argument" - ), - ) - - return supported, reason - - -# Each entry: (recipe_class_name, check_fn) -_FP8_RECIPE_CONFIGS = [ - ("DelayedScaling", fp8.check_fp8_support), - ("Float8CurrentScaling", fp8.check_fp8_support), - ("Float8BlockScaling", fp8.check_fp8_block_scaling_support), - ("MXFP8BlockScaling", fp8.check_mxfp8_support), - ("NVFP4BlockScaling", check_nvfp4_support), -] - - -def _parametrize_fp8_recipes(): - """Generate pytest.param objects with skip marks for unsupported FP8 recipes.""" - params = [] - for name, check_fn in _FP8_RECIPE_CONFIGS: - supported, reason = check_fn() - params.append( - pytest.param( - name, - id=name, - marks=pytest.mark.skipif(not supported, reason=reason), - ) - ) - return params - - -@pytest.fixture(params=_parametrize_fp8_recipes()) -def fp_recipe(request): - """Parametrized fixture providing FP8 recipe Hydra overrides for each supported TE recipe.""" - return request.param - - -def _run_test(fp_init, sharding_dims, recipe, layer_type): - test_path = Path(__file__).parent.resolve() / "run_fsdp2_model.py" - test_cmd = ["torchrun", f"--nproc_per_node={NUM_PROCS}", str(test_path)] - - if fp_init: - test_cmd += ["--fp8-init"] - - if len(sharding_dims) == 1: - test_cmd += ["--sharding-dims", str(sharding_dims[0])] - elif len(sharding_dims) == 2: - test_cmd += ["--sharding-dims", str(sharding_dims[0]), str(sharding_dims[1])] - else: - assert False - test_cmd += ["--recipe", recipe] - test_cmd += ["--layer-type", layer_type] - - subprocess.run(test_cmd, env=os.environ, check=True) +_FSDP2_DIR = Path(__file__).parent.resolve() / "fsdp2_tests" @pytest.mark.skipif(NUM_PROCS % 2 != 0, reason="Requires even number of GPUs") @pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") -@pytest.mark.parametrize("sharding_dims", ([NUM_PROCS], [2, NUM_PROCS // 2])) -@pytest.mark.parametrize("fp8_init", (False, True)) -@pytest.mark.parametrize("layer_type", ("LayerNormLinear", "TransformerLayer")) -def test_distributed(fp8_init, sharding_dims, fp_recipe, layer_type): - - if fp_recipe in ("Float8BlockScaling", "NVFP4BlockScaling") and fp8_init: - pytest.xfail(f"{fp_recipe} + fp8_init: test_fp8_fsdp2_allgather is currently failing.") - - _run_test(fp8_init, sharding_dims, fp_recipe, layer_type) - - -## ── FusedAdam + FSDP2 tests ───────────────────────────────────────── - - -def _run_fused_adam_test(test_name, recipe="delayed_scaling"): - """Launch an FSDP2 + FusedAdam test via torchrun.""" - test_path = Path(__file__).parent.resolve() / "run_fsdp2_fused_adam.py" - nproc = min(NUM_PROCS, 2) # These tests only need 2 GPUs - test_cmd = [ - "torchrun", - f"--nproc_per_node={nproc}", - str(test_path), - "--test", - test_name, - "--recipe", - recipe, - ] - - subprocess.run(test_cmd, env=os.environ, check=True) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_fp8_master_weights(fp_recipe): - """FusedAdam(master_weights=True) + FSDP2 + quantized_model_init (meta device init).""" - if fp_recipe in ("NVFP4BlockScaling",): - pytest.xfail( - f"{fp_recipe}: quantized_model_init and FSDP2 is not currently supported, since the " - "block tensor is dequantized before we flatten it for FSDP2." - ) - _run_fused_adam_test("fused_adam_fp8_master_weights", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_fp8_master_weights_no_meta(fp_recipe): - """FusedAdam(master_weights=True) + FSDP2 + quantized_model_init (CUDA init, no meta device). - - Block-scaling QuantizedTensors (MXFP8, Float8Blockwise, NVFP4) are wrapper - subclasses with data_ptr() == 0. Without meta-device init, FSDP2's - reset_sharded_param() crashes with 'invalid python storage'. - Per-tensor FP8 (DelayedScaling, Float8CurrentScaling) works because - Float8Tensor's storage is accessible. - """ - if fp_recipe in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): - pytest.xfail( - f"{fp_recipe}: FSDP2 without meta-device init crashes on block-scaling " - "QuantizedTensor wrapper subclasses (data_ptr() == 0). " - "Use device='meta' + reset_parameters() after sharding." - ) - _run_fused_adam_test("fused_adam_fp8_master_weights_no_meta", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_bf16(fp_recipe): - """FusedAdam(master_weights=True) + FSDP2 + bf16 params (no FP8).""" - _run_fused_adam_test("fused_adam_bf16", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_fp8_no_master(fp_recipe): - """FusedAdam(master_weights=False) + FSDP2 + FP8 params.""" - if fp_recipe in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): - pytest.xfail( - f"{fp_recipe}: FusedAdam without master_weights does not support " - "block-scaling quantized tensors. Use master_weights=True." - ) - _run_fused_adam_test("fused_adam_fp8_no_master", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_bf16_store_param_remainders(fp_recipe): - """FusedAdam(master_weights=True, store_param_remainders=True) + FSDP2 + bf16.""" - _run_fused_adam_test("fused_adam_bf16_store_param_remainders", fp_recipe) +def test_fsdp2_model_tests(): + """All FSDP2 model tests (parametrized internally by recipe, fp8_init, sharding, layer).""" + test_path = _FSDP2_DIR / "run_fsdp2_model.py" + result = subprocess.run( + [ + "torchrun", + f"--nproc_per_node={NUM_PROCS}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + ], + env=os.environ, + timeout=600, + ) + assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" @pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_dcp_output_parity(fp_recipe): - """DCP save/load round-trip into a fresh model produces identical outputs.""" - if fp_recipe == "MXFP8BlockScaling": - pytest.xfail( - "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access" - ) - - if fp_recipe == "NVFP4BlockScaling": - pytest.xfail( - "NVFP4BlockScaling: DCP load_state_dict triggers reset_sharded_param() " - "which calls data_ptr() on NVFP4Tensor wrapper subclass with invalid storage" - ) - - if fp_recipe == "Float8BlockScaling" and torch.cuda.get_device_capability()[0] == 12: - pytest.xfail( - "Float8BlockScaling is failing on SM120 with RuntimeError: " - "transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu:534 " - "in function quantize_transpose_vector_blockwise: Assertion failed: pow2_scale. On " - "Blackwell and newer, the FP8 block scaling recipe is emulated with MXFP8, which " - "requires using power of two scaling factors." - ) - - _run_fused_adam_test("dcp_output_parity", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_dcp_output_parity_async(fp_recipe): - """DCP save/load round-trip into a fresh model produces identical outputs.""" - if fp_recipe == "MXFP8BlockScaling": - pytest.xfail( - "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access: " - "/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh:92 in function " - "multi_tensor_apply: CUDA Error: an illegal memory access was encountered" - ) - - if fp_recipe == "NVFP4BlockScaling": - pytest.xfail( - "NVFP4BlockScaling: DCP load_state_dict triggers reset_sharded_param() " - "which calls data_ptr() on NVFP4Tensor wrapper subclass with invalid storage" - ) - - if fp_recipe == "Float8BlockScaling": - pytest.xfail( - "Float8BlockScaling: async DCP save/load round-trip produces different model " - "outputs — quantization metadata (scales) is not correctly persisted through " - "async distributed checkpointing. On SM120, additionally fails with pow2_scale " - "assertion in quantize_transpose_vector_blockwise." - ) - - _run_fused_adam_test("dcp_output_parity_async", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_safetensors_fp32_export(fp_recipe): - """Export FP32 model from optimizer master weights to safetensors.""" - if fp_recipe == "MXFP8BlockScaling": - pytest.xfail( - "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access" - ) - _run_fused_adam_test("safetensors_fp32_export", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -@pytest.mark.xfail( - reason=( - "fuse_wgrad_accumulation is incompatible with vanilla FSDP2: " - "autograd Function.apply unwraps DTensors to local tensors, so " - "main_grad (set on the DTensor) is inaccessible during backward. " - "Additionally, the fused wgrad GEMM bypasses FSDP2's reduce-scatter." - ), - raises=subprocess.CalledProcessError, - strict=True, -) -def test_fsdp2_fuse_wgrad_accumulation(fp_recipe): - """fuse_wgrad_accumulation=True + FSDP2 -- expected to fail.""" - _run_fused_adam_test("fuse_wgrad_accumulation", fp_recipe) +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_fsdp2_fused_adam_tests(): + """All FSDP2 FusedAdam tests (parametrized internally by recipe, test variant).""" + test_path = _FSDP2_DIR / "run_fsdp2_fused_adam.py" + nproc = min(NUM_PROCS, 2) + result = subprocess.run( + [ + "torchrun", + f"--nproc_per_node={nproc}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + ], + env=os.environ, + timeout=600, + ) + assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" def test_dummy() -> None: