diff --git a/.github/actions/cache-extensions/action.yml b/.github/actions/cache-extensions/action.yml index 7ddf65de8ce..a25dcb5fafc 100644 --- a/.github/actions/cache-extensions/action.yml +++ b/.github/actions/cache-extensions/action.yml @@ -1,22 +1,32 @@ name: Cache extensions -description: Cache and reuse JIT-compiled extensions (e.g. CUDA extensions) across runs. +description: Restore JIT-compiled extensions (e.g. CUDA extensions) cached by an earlier run. Pair with ./.github/actions/save-extensions. inputs: cache-key: description: Environment discriminator for the cache key (e.g. GPU arch + container image). required: true +outputs: + cache-hit: + description: Whether the extensions were restored from the exact key. + value: ${{ steps.cache.outputs.cache-hit }} + cache-key: + description: The key to pass to ./.github/actions/save-extensions. + value: ${{ steps.cache.outputs.cache-primary-key }} + runs: using: composite steps: - shell: bash run: echo "TORCH_EXTENSIONS_DIR=/root/.cache/torch_extensions" >> "$GITHUB_ENV" - id: cache - uses: actions/cache@v6 + uses: actions/cache/restore@v6 with: path: /root/.cache/torch_extensions - key: torch-ext-${{ inputs.cache-key }}-${{ hashFiles('modelopt/torch/kernels/quantization/**', 'modelopt/torch/quantization/extensions.py', 'modelopt/torch/utils/cpp_extension.py') - }} + # Only the JIT-built sources and their build scripts, so Triton-only edits keep the cache. + key: torch-ext-${{ inputs.cache-key }}-${{ hashFiles('modelopt/torch/kernels/quantization/**/*.cpp', 'modelopt/torch/kernels/quantization/**/*.cu', + 'modelopt/torch/kernels/quantization/**/*.cuh', 'modelopt/torch/kernels/quantization/**/*.h', 'modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py', + 'modelopt/torch/quantization/extensions.py', 'modelopt/torch/utils/cpp_extension.py') }} # On a cache hit, backdate sources below the cached objects so ninja reuses them (touching # the objects instead would desync ninja's .ninja_deps and force a rebuild). - if: steps.cache.outputs.cache-hit == 'true' diff --git a/.github/actions/save-extensions/action.yml b/.github/actions/save-extensions/action.yml new file mode 100644 index 00000000000..09f03d45794 --- /dev/null +++ b/.github/actions/save-extensions/action.yml @@ -0,0 +1,22 @@ +name: Save extensions +description: Save the JIT-compiled extensions restored by ./.github/actions/cache-extensions. Run it with `always()` so a failing or timed-out job still + populates the cache (actions/cache itself only saves when the job succeeds). + +inputs: + cache-key: + description: The `cache-key` output of ./.github/actions/cache-extensions. + required: true + +runs: + using: composite + steps: + # A job cancelled before building anything must not claim the key with an empty cache: + # later runs would restore it as an exact hit and never save again. + - id: built + shell: bash + run: find /root/.cache/torch_extensions -name '*.so' 2>/dev/null | grep -q . && echo "any=true" >> "$GITHUB_OUTPUT" || true + - if: steps.built.outputs.any == 'true' + uses: actions/cache/save@v6 + with: + path: /root/.cache/torch_extensions + key: ${{ inputs.cache-key }} diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml index e3378865b65..63147221b81 100644 --- a/.github/workflows/_example_tests_runner.yml +++ b/.github/workflows/_example_tests_runner.yml @@ -59,8 +59,10 @@ jobs: - uses: actions/checkout@v6 - uses: nv-gha-runners/setup-proxy-cache@main - uses: ./.github/actions/cache-extensions + id: ext-cache with: - cache-key: rtxpro6000-${{ inputs.docker_image }} + # Per example: jobs sharing an image build different extensions, and the first to save wins the key. + cache-key: rtxpro6000-example-${{ inputs.example }}-${{ inputs.docker_image }} - name: Setup environment variables run: | echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu:/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" >> $GITHUB_ENV @@ -113,6 +115,12 @@ jobs: pytest_args+=(-n 4) fi python3 -m pytest tests/examples/${{ inputs.example }} --cov "${pytest_args[@]}" + - name: Save extensions + # Not tied to the tests passing, see the action description + if: always() && steps.ext-cache.outputs.cache-hit != 'true' + uses: ./.github/actions/save-extensions + with: + cache-key: ${{ steps.ext-cache.outputs.cache-key }} - name: Flag allowed failure if: ${{ inputs.allow_failure && steps.run_tests.outcome == 'failure' }} run: | diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index 1c87931182e..67b98cbd230 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -131,7 +131,7 @@ jobs: contents: read secrets: inherit with: - docker_image: "nvcr.io/nvidia/pytorch:26.07-py3" + docker_image: "nvcr.io/nvidia/pytorch:26.09-py3" example: ${{ matrix.example }} timeout_minutes: 30 pip_install_extras: "[hf,dev-test]" @@ -172,6 +172,8 @@ jobs: contents: read secrets: inherit with: + # Newer releases break quickstart_multimodal.py (test_qwen_vl): rc21-rc23 import MultimodalConfig, + # which tensorrt_llm.llmapi does not export yet, and rc24+ import `fla`, which the image lacks. docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20" example: ${{ matrix.example }} pip_install_extras: "[hf,dev-test]" diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index b3b9bc6dfc7..b73f85affeb 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -45,25 +45,21 @@ jobs: include: - example: gpu timeout: 60 - container_image: nvcr.io/nvidia/pytorch:26.07-py3 + container_image: nvcr.io/nvidia/pytorch:26.09-py3 - example: gpu_megatron - timeout: 60 + timeout: 75 container_image: nvcr.io/nvidia/nemo:26.08 - example: gpu_trtllm timeout: 15 container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20 - # Covers the RoutedExperts layout (0.24+) - example: gpu_vllm - timeout: 15 + timeout: 30 container_image: docker.io/vllm/vllm-openai:v0.30.0 - # Covers the legacy FusedMoE / SharedFusedMoE registration branches (vLLM < 0.24) - - example: gpu_vllm - timeout: 15 - container_image: docker.io/vllm/vllm-openai:v0.20.0 runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }} timeout-minutes: ${{ matrix.timeout }} container: image: ${{ matrix.container_image }} + options: --shm-size=2gb # vLLM TP>1 tests on the 2-GPU runner need more than the default 64 MB env: GIT_DEPTH: 1000 # For correct version PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages @@ -79,8 +75,18 @@ jobs: - uses: actions/checkout@v6 - uses: nv-gha-runners/setup-proxy-cache@main - uses: ./.github/actions/cache-extensions + id: ext-cache with: - cache-key: rtxpro6000-${{ matrix.container_image }} + # Per job: jobs sharing an image build different extensions, and the first to save wins the key. + cache-key: rtxpro6000-${{ matrix.example }}-${{ matrix.container_image }} + - name: Restore wheels built from source + # mamba_ssm and fast-hadamard-transform take ~10 min to compile (see noxfile.py `gpu`) + id: wheelhouse + if: matrix.example == 'gpu' + uses: actions/cache/restore@v6 + with: + path: /root/.cache/wheelhouse + key: wheelhouse-${{ matrix.container_image }}-${{ hashFiles('noxfile.py') }} - name: Setup environment variables run: | echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu" >> $GITHUB_ENV @@ -88,9 +94,30 @@ jobs: env: COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml COVERAGE_FILE: ${{ github.workspace }}/.coverage + # Explicit: HOME is /github/home in container jobs, not the /root the cache steps use + MODELOPT_WHEELHOUSE: /root/.cache/wheelhouse run: | # Use `python3` (the vllm image has no `python` on PATH) python3 -m pip install nox && nox -s ${{ matrix.example }} + - name: Save extensions + # Not tied to the tests passing, see the action description + if: always() && steps.ext-cache.outputs.cache-hit != 'true' + uses: ./.github/actions/save-extensions + with: + cache-key: ${{ steps.ext-cache.outputs.cache-key }} + - name: Check the wheelhouse is complete + id: wheelhouse-complete + if: always() && matrix.example == 'gpu' + run: test -f /root/.cache/wheelhouse/.complete && echo "complete=true" >> "$GITHUB_OUTPUT" || true + - name: Save wheels built from source + # Not tied to the tests passing, but only once every wheel is built (see noxfile.py `gpu`). + if: >- + always() && matrix.example == 'gpu' && steps.wheelhouse.outputs.cache-hit != 'true' + && steps.wheelhouse-complete.outputs.complete == 'true' + uses: actions/cache/save@v6 + with: + path: /root/.cache/wheelhouse + key: wheelhouse-${{ matrix.container_image }}-${{ hashFiles('noxfile.py') }} - name: Upload GPU coverage to Codecov uses: codecov/codecov-action@v7 with: diff --git a/.github/workflows/regression_tests.yml b/.github/workflows/regression_tests.yml index d87ba8111c5..9c1154a2b65 100644 --- a/.github/workflows/regression_tests.yml +++ b/.github/workflows/regression_tests.yml @@ -43,7 +43,7 @@ jobs: runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }} timeout-minutes: 15 container: - image: nvcr.io/nvidia/pytorch:26.07-py3 + image: nvcr.io/nvidia/pytorch:26.09-py3 env: GIT_DEPTH: 1000 # For correct version PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index f605ddaf636..51f9392028b 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -72,7 +72,10 @@ from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper -from modelopt.torch.quantization.qtensor.nvfp4_tensor import _cast_per_block_scale_to_fp8 +from modelopt.torch.quantization.qtensor.nvfp4_tensor import ( + _cast_per_block_scale_to_fp8, + cutlass_fp4_scale_to_modelopt_fp4_scale, +) from modelopt.torch.quantization.utils import ( fsdp2_aware_weight_update, module_name_maps, @@ -591,15 +594,6 @@ def _compressed_per_block_scale( scale = getattr(weight_quantizer, "_scale", None) if scale is None or not (scale.dtype == torch.uint8 and scale.ndim == 1): return scale - try: - from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import ( - cutlass_fp4_scale_to_modelopt_fp4_scale, - ) - except ImportError as e: - raise ImportError( - "This weight was compressed by TensorRT-LLM, so its NVFP4 block scale is " - "cutlass-swizzled, but tensorrt_llm cannot be imported to convert it for export." - ) from e return cutlass_fp4_scale_to_modelopt_fp4_scale(scale, weight.metadata["shape"][-2:]) diff --git a/modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py b/modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py index 713b5f82f9b..8c6d3dfacbf 100644 --- a/modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py +++ b/modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py @@ -59,7 +59,8 @@ def _get_cuda_module(): "--use_fast_math", "-lineinfo", "--ptxas-options=-v", - "-std=c++17", + # No -std flag: torch's load() adds the C++ standard its headers need (C++20 in + # recent releases), and a hardcoded -std=c++17 fails to build there. ], ) return _cuda_module diff --git a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py index a56588b21d8..c9ee0aa24bd 100644 --- a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py +++ b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py @@ -24,8 +24,6 @@ the per-block scale is the identity, so the kernel can use ``scale = candidate * global_amax / 6.0`` without an explicit FP8 cast — making it runnable on any CUDA GPU with Triton (no ``tl.float8e4nv`` requirement). - -Tile shape (``BLOCKS_PER_PROGRAM``) and ``num_warps`` are autotuned per ``N_BLOCKS``. """ import torch @@ -43,18 +41,13 @@ ] -# Selected from a (BLOCKS_PER_PROGRAM, num_warps) sweep on B300: -# BPP=16,nw=2: 6.06 ms BPP=32,nw=4: 6.06 ms BPP=64,nw=8: 5.08 ms -# The smaller-tile entries cover cases where N_BLOCKS is small enough that BPP=64 -# would underfill the SMs. -_FP8_SWEEP_AUTOTUNE_CONFIGS = [ - triton.Config({"BLOCKS_PER_PROGRAM": 16}, num_warps=2), - triton.Config({"BLOCKS_PER_PROGRAM": 32}, num_warps=4), - triton.Config({"BLOCKS_PER_PROGRAM": 64}, num_warps=8), -] +# Fastest in a (BLOCKS_PER_PROGRAM, num_warps) sweep on B300 (BPP=16/32: 6.06 ms, BPP=64: 5.08 ms) +# and within noise of the others on small shapes. Hard-coded since autotuning compiles every +# candidate (~10 s each) on first use. +_FP8_SWEEP_BLOCKS_PER_PROGRAM = 64 +_FP8_SWEEP_NUM_WARPS = 8 -@triton.autotune(configs=_FP8_SWEEP_AUTOTUNE_CONFIGS, key=["N_BLOCKS"]) @triton.jit def _fp8_scale_sweep_kernel( x_ptr, # [N_BLOCKS * BLOCK_SIZE], any float dtype (loaded as fp32) @@ -148,7 +141,7 @@ def nvfp4_fp8_scale_sweep( candidates = fp8_scale_candidates(x.device).to(dtype=torch.float32) global_amax_f32 = global_amax.detach().to(device=x.device, dtype=torch.float32).reshape(1) - grid = lambda meta: (triton.cdiv(n_blocks, meta["BLOCKS_PER_PROGRAM"]),) + grid = (triton.cdiv(n_blocks, _FP8_SWEEP_BLOCKS_PER_PROGRAM),) with torch.cuda.device(x.device): _fp8_scale_sweep_kernel[grid]( x_flat, @@ -158,6 +151,8 @@ def nvfp4_fp8_scale_sweep( n_blocks, BLOCK_SIZE=block_size, NUM_CANDIDATES=int(candidates.numel()), + BLOCKS_PER_PROGRAM=_FP8_SWEEP_BLOCKS_PER_PROGRAM, + num_warps=_FP8_SWEEP_NUM_WARPS, ) return best_amax diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index b3095e0ca98..8726dc5c719 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -46,6 +46,24 @@ def _cast_per_block_scale_to_fp8( return per_block_scale.clamp(min=2**-9, max=E4M3_MAX).to(torch.float8_e4m3fn) +def cutlass_fp4_scale_to_modelopt_fp4_scale( + cutlass_scale: torch.Tensor, weight_shape: tuple[int, int] +) -> torch.Tensor: + """Undo the cuBLAS/cutlass tiled swizzle of NVFP4 block scales (TensorRT-LLM ``fp4_quantize``). + + The swizzled scale holds 128x4 tiles, each stored as ``[row % 32, row // 32, col]``; this is the + inverse of the rearrangement in ``swizzle_nvfp4_scales``. Returns the ``[m, n // 16]`` E4M3 + scales of an ``[m, n]`` weight with block size 16. + """ + m, n = weight_shape + n //= 16 + row_blocks, col_blocks = -(-m // 128), -(-n // 4) + tiles = cutlass_scale.reshape(row_blocks, col_blocks, 32, 4, 4) + padded = tiles.permute(0, 3, 2, 1, 4).reshape(row_blocks * 128, col_blocks * 4) + # contiguous(): the slice is a strided view, and fp4_dequantize indexes the scale flatly. + return padded[:m, :n].contiguous().view(torch.float8_e4m3fn) + + class NVFP4QTensor(BaseQuantizedTensor): """Implements the NVFP4 quantization on tensors for more efficient storage or computation. @@ -294,11 +312,6 @@ def quantize( try: import tensorrt_llm # noqa: F401 - # Make sure this utils is available for dequantize - from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import ( - cutlass_fp4_scale_to_modelopt_fp4_scale, # noqa: F401 - ) - packed_weight, weights_scaling_factor = torch.ops.trtllm.fp4_quantize( input, 1.0 / weights_scaling_factor_2, block_size, False ) @@ -364,18 +377,9 @@ def _unpack_tensor(input: torch.Tensor): # Get scales from kwargs if kwarg["scale"].dtype == torch.uint8 and kwarg["scale"].ndim == 1: # If quantization is done by trtllm, convert cutlass fp4 scale to modelopt fp4 scale - try: - from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import ( - cutlass_fp4_scale_to_modelopt_fp4_scale, - ) - - kwarg["scale"] = cutlass_fp4_scale_to_modelopt_fp4_scale( - kwarg["scale"], self.metadata["shape"][-2:] - ) - except ImportError as e: - raise ImportError( - "This tensor is quantized by trtllm, but tensorrt_llm cannot be imported." - ) from e + kwarg["scale"] = cutlass_fp4_scale_to_modelopt_fp4_scale( + kwarg["scale"], self.metadata["shape"][-2:] + ) if fast: from modelopt.torch.kernels.quantization.gemm.fp4_kernel import fp4_dequantize diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 63f349dd10a..c0f3884f41d 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -828,20 +828,37 @@ def _init_mp_dtypes(self) -> None: @contextmanager def _disable_fsdp_unshard_reshard(layer): - """Disable FSDP2 unshard/reshard if *layer* is FSDP-wrapped.""" - if isinstance(layer, FSDPModule): - _pg_cls = torch.distributed.fsdp._fully_shard._fsdp_param_group.FSDPParamGroup - orig_unshard = _pg_cls.unshard - orig_reshard = _pg_cls.reshard - _pg_cls.unshard = lambda self, async_op=False: None - _pg_cls.reshard = lambda self: None - try: - yield - finally: - _pg_cls.unshard = orig_unshard - _pg_cls.reshard = orig_reshard - else: + """Make the FSDP2 param groups inside *layer* skip unshard/reshard (already gathered). + + Patched per group instance, not on ``FSDPParamGroup``: the rest of the model still runs forward + (e.g. while collecting the next layer's inputs) and must unshard as usual. + """ + if not isinstance(layer, FSDPModule): # nothing was gathered for its FSDP submodules + yield + return + groups = { + id(group): group # ``fully_shard([a, b])`` gives several modules one group + for module in layer.modules() + if isinstance(module, FSDPModule) + and (group := fully_shard.state(module)._fsdp_param_group) is not None + }.values() + # Contexts nest (``persistent_materialization(writeback=False)`` enters this twice), so restore + # whatever instance attributes were there before instead of deleting them. + saved = [ + (group, {k: group.__dict__.get(k) for k in ("unshard", "reshard")}) for group in groups + ] + for group in groups: + group.unshard = lambda async_op=False: None + group.reshard = lambda: None + try: yield + finally: + for group, prior in saved: + for name, value in prior.items(): + if value is None: + del group.__dict__[name] + else: + setattr(group, name, value) def build_param_index(model): diff --git a/modelopt/torch/speculative/plugins/hf_eagle.py b/modelopt/torch/speculative/plugins/hf_eagle.py index e7c931c1c86..910d5907870 100644 --- a/modelopt/torch/speculative/plugins/hf_eagle.py +++ b/modelopt/torch/speculative/plugins/hf_eagle.py @@ -147,7 +147,9 @@ def _activate_torch_compile(self): compile_targets: list[tuple[str, dict[str, Any]]] = [ ("_prepare_eagle_inputs", {}), - ("_eagle_forward", {"mode": "max-autotune"}), + # No CUDA graphs: "max-autotune" enables them, and recording them intermittently fails + # in training (cudaStreamCaptureStatusInvalidated / "graph recording observed ..."). + ("_eagle_forward", {"mode": "max-autotune-no-cudagraphs"}), ("_eagle_loss", {"fullgraph": True}), ] for name, kwargs in compile_targets: @@ -752,7 +754,8 @@ def forward( base_outputs, ) - # Init rotary_emb outside _eagle_forward so inv_freq is not captured by CUDAGraph. + # Init rotary_emb outside _eagle_forward so inv_freq is not captured by CUDAGraph. Capture is + # off for now (see _activate_torch_compile); this keeps re-enabling it safe. self.eagle_module._maybe_init_rope(device=eagle_input_hiddens.device) # ====Run eagle forward with extra training-time-test steps==== @@ -900,7 +903,8 @@ def pseudo_speculative_generate( else: eagle_input_hidden_states = base_model_hidden_states - # Init rotary_emb outside _eagle_forward so inv_freq is not captured by CUDAGraph. + # Init rotary_emb outside _eagle_forward so inv_freq is not captured by CUDAGraph. Capture is + # off for now (see _activate_torch_compile); this keeps re-enabling it safe. self.eagle_module._maybe_init_rope(device=eagle_input_hidden_states.device) draft_tokens = [] diff --git a/noxfile.py b/noxfile.py index 27d65b26810..321a8924210 100644 --- a/noxfile.py +++ b/noxfile.py @@ -26,6 +26,8 @@ import glob import os import shutil +import tarfile +import tempfile import nox @@ -76,7 +78,7 @@ def unit(session, torch_ver, tf_ver): ) if run_manual: # A separate process: some manual tests (e.g. multi-threaded from_pretrained) leak global - # state that breaks later tests. No coverage, so it cannot overwrite the main run's data. + # state that breaks later tests. No --cov here, so this run writes no coverage data. session.run( "python", # Caller args first: the last ``-m`` wins, and it must be ours. @@ -113,30 +115,88 @@ def partial_unit(session, subset): # existing Python environment (e.g. /opt/venv in NeMo) instead of an isolated one. # Use `python -m pip/pytest` to ensure the container's active venv Python is used, # not a stale PATH entry (e.g. NeMo container has pip → /usr/local/bin/pip but python → /opt/venv/bin/python). +# CUDA packages compiled against the container's torch, which takes minutes each. Pinned so the CI +# cache of the built wheels (keyed on this file and the container image, see gpu_tests.yml) stays +# valid: bump a pin deliberately, which also invalidates that cache. +_SOURCE_BUILT_WHEELS = { + "fast_hadamard_transform": "git+https://github.com/Dao-AILab/fast-hadamard-transform.git@e7706faf8d1c3b9f241e36860640ad1dac644ede", + "mamba_ssm": "mamba_ssm==2.3.2.post1", + "causal_conv1d": "causal-conv1d==1.7.0", +} + + +# Their setup.py files hardcode -std=c++17, but the torch headers of recent containers need C++20. +_BUILD_AS_CXX20 = ("mamba_ssm", "causal_conv1d") +# Without these, setup.py may download a matching prebuilt wheel instead of compiling the patched source. +_FORCE_BUILD = { + "mamba_ssm": {"MAMBA_FORCE_BUILD": "TRUE"}, + "causal_conv1d": {"CAUSAL_CONV1D_FORCE_BUILD": "TRUE"}, +} + + +def _wheelhouse(): + return os.environ.get("MODELOPT_WHEELHOUSE", os.path.expanduser("~/.cache/wheelhouse")) + + +def _build_wheel(session, name, wheelhouse): + pip_wheel = ("python", "-m", "pip", "wheel", "--no-build-isolation", "--no-deps") + pip_wheel += ("--wheel-dir", wheelhouse) + requirement = _SOURCE_BUILT_WHEELS[name] + if name not in _BUILD_AS_CXX20: + session.run(*pip_wheel, requirement) + return + with tempfile.TemporaryDirectory() as tmp: + session.run( + *("python", "-m", "pip", "download", "--no-deps", "--no-binary", ":all:"), + *("--no-build-isolation", "--dest", tmp, requirement), + ) + sdist = glob.glob(os.path.join(tmp, "*.tar.gz"))[0] + with tarfile.open(sdist) as tar: + tar.extractall(tmp, filter="data") + source = sdist.removesuffix(".tar.gz") + setup_py = os.path.join(source, "setup.py") + with open(setup_py) as f: + patched = f.read().replace("-std=c++17", "-std=c++20") + with open(setup_py, "w") as f: + f.write(patched) + session.run(*pip_wheel, source, env=_FORCE_BUILD[name]) + + +def _find_wheel(wheelhouse, name): + # ``-[0-9]`` so ``causal_conv1d`` does not also match a ``causal_conv1d_cuda-*`` wheel. + found = sorted(glob.glob(os.path.join(wheelhouse, f"{name}-[0-9]*.whl"))) + return found[-1] if found else None + + +def _install_built_from_source(session, *names): + """Install the named ``_SOURCE_BUILT_WHEELS``, building a wheel only if the wheelhouse lacks it. + + CI caches the wheelhouse (``MODELOPT_WHEELHOUSE``) per container image, so only the first run on + an image builds. + """ + wheelhouse = _wheelhouse() + wheels = [] + for name in names: + wheel = _find_wheel(wheelhouse, name) + if wheel is None: + _build_wheel(session, name, wheelhouse) + wheel = _find_wheel(wheelhouse, name) + if wheel is None: + raise RuntimeError(f"No wheel for {name!r} in {wheelhouse} after building it") + wheels.append(wheel) + session.run("python", "-m", "pip", "install", *wheels) + + # Container: nvcr.io/nvidia/pytorch:26.07-py3 or later with CUDA 13 @nox.session(venv_backend="none") def gpu(session): - # tests/gpu/_extensions/test_onnx_extensions.py fails for newer containers - # until https://github.com/tbenthompson/cppimport/pull/98 - session.run( - "python", - "-m", - "pip", - "install", - "--no-build-isolation", - "git+https://github.com/Dao-AILab/fast-hadamard-transform.git", - ) + _install_built_from_source(session, "fast_hadamard_transform") session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test,dev-fla]") - session.run( - "python", - "-m", - "pip", - "install", - "--no-build-isolation", - # Install the latest *released* sdists (built against the container torch) - "mamba_ssm", - "causal-conv1d", - ) + _install_built_from_source(session, "mamba_ssm", "causal_conv1d") + # Every wheel is built: CI saves the wheelhouse only then, so a run that stopped halfway cannot + # freeze a partial wheelhouse under the cache key. + with open(os.path.join(_wheelhouse(), ".complete"), "w"): + pass session.run("python", "-m", "pytest", "tests/gpu", *_cov_args()) diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index ea954efbcf5..02d5f6e8869 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -187,21 +187,27 @@ def _worker_loop(rank, world_size, backend, port, cmd_queue, result_queue, teard fn, args, kwargs = cmd status = "ok" tb = None + reported = False try: fn(rank, world_size, *args, **kwargs) except Exception: status = "error" tb = traceback.format_exc() + # Report before teardown: it may block on a collective whose peer is still stuck in + # ``fn``, and the parent's failure grace period starts when it hears of a failure. + result_queue.put((status, rank, tb)) + reported = True finally: if teardown_fn is not None: try: teardown_fn(rank, world_size) except Exception as e: - print(f"Error tearing down worker: {e}") - status = "error" - teardown_tb = traceback.format_exc() - tb = (tb + "\n" if tb else "") + f"[teardown] {teardown_tb}" - result_queue.put((status, rank, tb)) + print(f"Error tearing down worker: {e}\n{traceback.format_exc()}") + if not reported: # else the job failure is already sent; the teardown error is printed above + status = "error" + tb = f"[teardown] {traceback.format_exc()}" + if not reported: + result_queue.put((status, rank, tb)) dist.destroy_process_group() @@ -211,11 +217,11 @@ def run(self, fn, *args, **kwargs): ``fn`` is called as ``fn(rank, world_size, *args, **kwargs)`` and must be picklable (top-level function or ``functools.partial`` of one). - If the wait is interrupted (e.g. by ``pytest-timeout``), or a rank fails and its peers do - not finish within ``failure_grace`` (they are typically blocked in a collective waiting for - it), the workers are still inside the job and every later ``run`` would hang until its own - timeout. The workers are killed instead and the pool is marked ``broken``; callers that - want to continue need a new pool. + If the job fails or the wait is interrupted (e.g. by ``pytest-timeout``), the workers may + still be inside the job or its teardown, typically blocked in a collective, and every later + ``run`` would hang until its own timeout. The workers are killed instead and the pool is + marked ``broken``; callers that want to continue need a new pool (``get_live_pool``). After + a rank fails, its peers get ``failure_grace`` seconds to report before the kill. """ if self.broken: raise RuntimeError("Worker pool was killed after an earlier job hung; create a new one") @@ -240,12 +246,16 @@ def run(self, fn, *args, **kwargs): raise if errors: + if not self.broken: + # Every rank replied, but a failed rank reports before its teardown, which can still + # be blocked; never hand these workers to another job. + self._abort(dump_stacks=False) raise RuntimeError("Worker(s) failed:\n" + "\n".join(errors)) - def _abort(self): - """Print every worker's stack (to stderr, for the CI log), then kill the workers.""" + def _abort(self, dump_stacks=True): + """Kill the workers, first printing their stacks (to stderr, for the CI log) if asked.""" self.broken = True - if hasattr(signal, "SIGUSR1"): + if dump_stacks and hasattr(signal, "SIGUSR1"): for p in self._processes: with contextlib.suppress(OSError): os.kill(p.pid, signal.SIGUSR1) diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py index 6977061d282..3bdcc460c27 100644 --- a/tests/examples/megatron_bridge/test_distill.py +++ b/tests/examples/megatron_bridge/test_distill.py @@ -130,7 +130,11 @@ def test_distill_llm_sft_hf_chat(tmp_path, num_gpus, packed): ``packed`` runs an MoE student with conversations packed per micro-batch and the aux loss off. """ create_dir = create_tiny_qwen3_moe_dir if packed else create_tiny_qwen3_dir - teacher_hf_path = create_dir(tmp_path, with_tokenizer=True) + # The default head_dim=2 leaves only TE's unfused attention, which the 2-GPU sm_120 runner + # rejects for packed (THD) input; head_dim=8 also qualifies the flash and cuDNN backends. + teacher_hf_path = create_dir( + tmp_path, with_tokenizer=True, **({"head_dim": 8} if packed else {}) + ) train_iters = 2 gbs = 4 # More records than train_iters * gbs so the sampler does not run dry; every 4th is multi-turn. diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py index 4953db2c737..222054ea187 100644 --- a/tests/gpu/torch/puzzletron/test_puzzletron.py +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -32,19 +32,25 @@ # SEED = 1234 +_MANUAL = pytest.mark.manual(reason="~40s each; run with --run-manual") @pytest.mark.parametrize( ("hf_model_name", "converter", "hybrid_override_pattern", "has_moe_layers"), [ - ("meta-llama/Llama-3.1-8B-Instruct", "llama", None, False), - ("meta-llama/Llama-3.2-3B-Instruct", "llama", None, False), - ("mistralai/Mistral-Small-24B-Instruct-2501", "mistral_small", None, False), + # Puzzletron is not actively maintained, so CI runs only Nemotron-3-Nano, gpt-oss and Qwen3-VL. + pytest.param("meta-llama/Llama-3.1-8B-Instruct", "llama", None, False, marks=_MANUAL), + pytest.param("meta-llama/Llama-3.2-3B-Instruct", "llama", None, False, marks=_MANUAL), + pytest.param( + "mistralai/Mistral-Small-24B-Instruct-2501", "mistral_small", None, False, marks=_MANUAL + ), ("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16", "nemotron_h", "*E", True), - ("nvidia/NVIDIA-Nemotron-Nano-12B-v2", "nemotron_h_v2", "*-", False), + pytest.param( + "nvidia/NVIDIA-Nemotron-Nano-12B-v2", "nemotron_h_v2", "*-", False, marks=_MANUAL + ), ("openai/gpt-oss-20b", "gpt_oss", None, True), - ("Qwen/Qwen2.5-7B-Instruct", "qwen2", None, False), - ("Qwen/Qwen3-8B", "qwen3", None, False), + pytest.param("Qwen/Qwen2.5-7B-Instruct", "qwen2", None, False, marks=_MANUAL), + pytest.param("Qwen/Qwen3-8B", "qwen3", None, False, marks=_MANUAL), ("Qwen/Qwen3-VL-30B-A3B-Instruct", "qwen3_vl", None, True), ], ) diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py index 69f0a117b63..078bf3f6b5a 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py @@ -287,7 +287,7 @@ def sharded_forward(*args, input_ids=None, **kwargs): model._base_model.forward = sharded_forward def test_generate_gathers_sharded_hidden_states(self): - # 8 layers so the default target_layer_ids ([1, 5]) straddle the split; 4 gives [1, 1]. + # 8 layers so the default target_layer_ids land in different halves; 4 collapses them. model = get_tiny_llama(num_hidden_layers=8) mtsp.convert(model, [("dflash", get_dflash_config(block_size=4))]) model = model.to("cuda:1").eval() diff --git a/tests/gpu/torch/utils/test_distributed_worker_pool.py b/tests/gpu/torch/utils/test_distributed_worker_pool.py index 7ae3a966278..36dce35968f 100644 --- a/tests/gpu/torch/utils/test_distributed_worker_pool.py +++ b/tests/gpu/torch/utils/test_distributed_worker_pool.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -17,21 +17,28 @@ import queue import time -from types import SimpleNamespace import pytest import torch.distributed as dist -from _test_utils.torch.distributed.utils import DistributedWorkerPool, get_live_pool +from _test_utils.torch.distributed.utils import DistributedWorkerPool def _hang(rank, size): time.sleep(600) -def _rank_0_fails_and_rank_1_waits_for_it(rank, size): +def _rank_0_fails_and_rank_1_hangs(rank, size): if rank == 0: raise ValueError("rank 0 failed") - dist.barrier() # never completes: its partner left + time.sleep(600) + + +def _fail(rank, size): + raise ValueError("job failed") + + +def _teardown_barrier(rank, size): + dist.barrier() # like megatron_worker_teardown: blocks while a peer is still in the job def test_a_hung_job_kills_the_pool_instead_of_stalling_later_runs(capfd): @@ -52,30 +59,27 @@ def test_a_hung_job_kills_the_pool_instead_of_stalling_later_runs(capfd): pool.shutdown() -def test_a_failed_rank_does_not_make_its_peers_wait_out_the_job_timeout(): - pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None) +def test_a_failed_rank_is_reported_even_when_its_teardown_blocks_on_a_hung_peer(): + pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=_teardown_barrier) try: pool.failure_grace = 2 + pool.job_timeout = 120 # without the early report the parent would wait this long start = time.monotonic() with pytest.raises(RuntimeError, match="rank 0 failed"): - pool.run(_rank_0_fails_and_rank_1_waits_for_it) + pool.run(_rank_0_fails_and_rank_1_hangs) - assert time.monotonic() - start < 60 # not the 600 s job timeout + assert time.monotonic() - start < 25 assert pool.broken finally: pool.shutdown() -def test_a_killed_pool_is_replaced_and_a_live_one_is_reused(): - pools = {} - made = [] - - def make(): - made.append(SimpleNamespace(broken=False)) - return made[-1] - - first = get_live_pool(pools, 2, make) - assert get_live_pool(pools, 2, make) is first - first.broken = True - second = get_live_pool(pools, 2, make) - assert second is not first and len(made) == 2 +def test_a_failed_job_retires_the_pool_even_when_every_rank_replied(): + """A failed rank reports before its teardown, which may still be blocked; never reuse it.""" + pool = DistributedWorkerPool(world_size=1, backend="gloo", teardown_fn=None) + try: + with pytest.raises(RuntimeError, match="job failed"): + pool.run(_fail) + assert pool.broken + finally: + pool.shutdown() diff --git a/tests/gpu_megatron/_extensions/test_torch_extensions.py b/tests/gpu_megatron/_extensions/test_torch_extensions.py deleted file mode 120000 index a07099dfd85..00000000000 --- a/tests/gpu_megatron/_extensions/test_torch_extensions.py +++ /dev/null @@ -1 +0,0 @@ -../../gpu/_extensions/test_torch_extensions.py \ No newline at end of file diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index 5760b9fe22d..51261e30023 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -47,10 +47,10 @@ def _prebuild_quant_cuda_extensions(): extensions during session setup and let tests fall back to on-demand JIT if needed. Doing it here in session setup (``pyproject`` sets ``timeout_func_only``) keeps the - build off the per-test clock and, unlike the ``_extensions/test_torch_extensions.py`` - prebuild tests, runs regardless of test selection/ordering (e.g. ``-k`` filters) and - is not itself capped by a per-test timeout. Worker subprocesses then load the cached - .so from the shared ``TORCH_EXTENSIONS_DIR``. + build off the per-test clock, runs regardless of test selection/ordering (e.g. ``-k`` + filters) and is not itself capped by a per-test timeout. Worker subprocesses then load + the cached .so from the shared ``TORCH_EXTENSIONS_DIR``. (The extension tests themselves + live in ``tests/gpu/_extensions``.) """ ext.precompile() @@ -99,7 +99,9 @@ def _pool_cache(): instead of spawning duplicates — e.g. on a 2-GPU runner ``dist_workers`` and ``dist_workers_size_2`` are both size 2. The cache is module-scoped and torn down at module end, so workers are never reused across modules (avoids cross-test - state contamination). + state contamination). The ``dist_workers*`` fixtures are deliberately function-scoped so + ``get_live_pool`` can swap in a new pool after one was killed; they therefore cannot be + requested from a module- or session-scoped fixture. """ pools: dict[int, DistributedWorkerPool] = {} yield pools diff --git a/tests/gpu_megatron/torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py b/tests/gpu_megatron/torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py index 8c98f7127fa..55404619de9 100644 --- a/tests/gpu_megatron/torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py +++ b/tests/gpu_megatron/torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py @@ -14,6 +14,7 @@ # limitations under the License. +import pytest import torch from _test_utils.import_helper import skip_if_no_mamba @@ -133,6 +134,7 @@ def _test_mamba_search_space(rank, size): assert not any(named_dynamic_modules(model)) +@pytest.mark.timeout(150) # usually ~70s, but up to 123s on slow runners def test_mamba_search_space(dist_workers): dist_workers.run(_test_mamba_search_space) diff --git a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py index 5dd0815080c..b676e1e3f85 100644 --- a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py +++ b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py @@ -40,15 +40,6 @@ def test_worker_context_falls_back_to_spawn_without_forkserver(monkeypatch): _prestart_forkserver() -# Depends on the external HF datasets-server, which intermittently returns 5xx -# (e.g. 503) on the /splits lookup. Allow-fail the transient outage while keeping -# real assertion failures visible (raises=RuntimeError only, strict=False so a -# normal pass doesn't xpass-fail). download_hf_dataset_as_jsonl already retries. -@pytest.mark.xfail( - raises=RuntimeError, - strict=False, - reason="Flaky: transient HF datasets-server 5xx on the /splits lookup", -) def test_megatron_preprocess_data_with_jsonl_path(tmp_path): input_jsonl = tmp_path / "minipile.jsonl" input_jsonl.write_text( @@ -109,12 +100,6 @@ def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path): assert limited_size < full_size -# Allow-fail transient HF datasets-server 5xx (see the /splits note above). -@pytest.mark.xfail( - raises=RuntimeError, - strict=False, - reason="Flaky: transient HF datasets-server 5xx on the /splits lookup", -) def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path): common_args = { "hf_dataset": "nanotron/minipile_100_samples", @@ -141,12 +126,6 @@ def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path): assert limited_size < full_size -# Allow-fail transient HF datasets-server 5xx (see the /splits note above). -@pytest.mark.xfail( - raises=RuntimeError, - strict=False, - reason="Flaky: transient HF datasets-server 5xx on the /splits lookup", -) def test_megatron_preprocess_data_hf_split_resume_uses_cached_token_count(tmp_path): args = { "hf_dataset": "nanotron/minipile_100_samples", @@ -168,12 +147,6 @@ def test_megatron_preprocess_data_hf_split_resume_uses_cached_token_count(tmp_pa assert set(tmp_path.iterdir()) == cached_files -# Allow-fail transient HF datasets-server 5xx (see the /splits note above). -@pytest.mark.xfail( - raises=RuntimeError, - strict=False, - reason="Flaky: transient HF datasets-server 5xx on the /splits lookup", -) @pytest.mark.parametrize( ("hf_dataset", "hf_split", "json_keys"), [ @@ -318,12 +291,6 @@ def test_megatron_preprocess_data_tool_calls_arguments_normalized(tmp_path): ) -# Allow-fail transient HF datasets-server 5xx (see the /splits note above). -@pytest.mark.xfail( - raises=RuntimeError, - strict=False, - reason="Flaky: transient HF datasets-server 5xx on the /splits lookup", -) def test_megatron_preprocess_data_hf_streaming_warning(tmp_path): # hf_streaming without hf_max_samples_per_split should warn and fall back to non-streaming with pytest.warns(UserWarning, match="hf_streaming"): diff --git a/tests/gpu_trtllm/torch/quantization/backends/test_nvfp4_gemm.py b/tests/gpu_trtllm/torch/quantization/backends/test_nvfp4_gemm.py index afa0aa1f2de..38c22ec3533 100644 --- a/tests/gpu_trtllm/torch/quantization/backends/test_nvfp4_gemm.py +++ b/tests/gpu_trtllm/torch/quantization/backends/test_nvfp4_gemm.py @@ -15,13 +15,11 @@ import pytest import torch from _test_utils.torch.quantization.models import OneLayerLinear -from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import ( - cutlass_fp4_scale_to_modelopt_fp4_scale, -) import modelopt.torch.quantization as mtq from modelopt.torch.quantization.backends.utils import fp4_compatible from modelopt.torch.quantization.qtensor import NVFP4QTensor +from modelopt.torch.quantization.qtensor.nvfp4_tensor import cutlass_fp4_scale_to_modelopt_fp4_scale @pytest.mark.skipif(not fp4_compatible(), reason="FP4 is not supported on this GPU") diff --git a/tests/unit/torch/quantization/test_nvfp4_tensor.py b/tests/unit/torch/quantization/test_nvfp4_tensor.py index a2a977f45b1..a7b64515e3a 100644 --- a/tests/unit/torch/quantization/test_nvfp4_tensor.py +++ b/tests/unit/torch/quantization/test_nvfp4_tensor.py @@ -24,6 +24,7 @@ from modelopt.torch.quantization.qtensor.nvfp4_tensor import ( NVFP4QTensor, _cast_per_block_scale_to_fp8, + cutlass_fp4_scale_to_modelopt_fp4_scale, ) _FP8_E4M3FN_MIN = 2**-9 # 0.001953125 — smallest positive FP8 E4M3FN subnormal @@ -127,3 +128,25 @@ def test_static_path_no_nan_when_block_amax_zero(self): assert (per_block_scale_f32 <= _FP8_E4M3FN_MAX).all(), ( f"Static per-block scale exceeds FP8 max 448: {per_block_scale_f32.tolist()}" ) + + +def _cublas_swizzle(scale): + """Flat cuBLAS tiled layout of a ``[rows, cols]`` block-scale matrix (the 128x4-tile swizzle).""" + rows, cols = scale.shape + row_blocks, col_blocks = -(-rows // 128), -(-cols // 4) + padded = torch.zeros(row_blocks * 128, col_blocks * 4, dtype=scale.dtype) + padded[:rows, :cols] = scale + blocks = padded.view(row_blocks, 128, col_blocks, 4).permute(0, 2, 1, 3) + return blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1) + + +@pytest.mark.parametrize("weight_shape", [(128, 64), (3, 16), (130, 80), (256, 128)]) +def test_cutlass_fp4_scale_to_modelopt_fp4_scale_undoes_the_swizzle(weight_shape): + m, n = weight_shape + scale = torch.randint(0, 256, (m, n // 16), dtype=torch.uint8) + + restored = cutlass_fp4_scale_to_modelopt_fp4_scale(_cublas_swizzle(scale), weight_shape) + + assert restored.dtype == torch.float8_e4m3fn + assert restored.is_contiguous() # consumers index it flatly + assert torch.equal(restored.view(torch.uint8), scale) diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py index 933c553ce12..0caab634562 100644 --- a/tests/unit/torch/quantization/test_utils.py +++ b/tests/unit/torch/quantization/test_utils.py @@ -15,12 +15,18 @@ import pytest import torch +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard from modelopt.torch.quantization.utils import ( convert_quantization_axis_to_reduce_axis, reduce_amax, reduce_block_amax, ) +from modelopt.torch.quantization.utils.core_utils import ( + _disable_fsdp_unshard_reshard, + persistent_materialization, +) from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @@ -333,3 +339,49 @@ def _forward_loop(m): assert "L1" not in call_log_for_layer1 assert torch.allclose(inp2[0][0][0], torch.tensor([[10.0 + 1.0 + 2.0]])) + + +def _forward_through_other_layers_while_one_is_materialized(rank, size, writeback): + mesh = init_device_mesh("cpu", (size,)) + torch.manual_seed(0) + model = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8)) + for layer in model: + fully_shard(layer, mesh=mesh) + fully_shard(model, mesh=mesh) + x = torch.randn(2, 8) + expected = model(x) + # Layerwise calibration runs the whole model to collect the next layer's inputs while one layer + # is materialized; the other FSDP modules must still unshard. + with persistent_materialization(model[0], writeback=writeback): + torch.testing.assert_close(model(x), expected) + torch.testing.assert_close(model(x), expected) # FSDP unshard/reshard restored on exit + + +# writeback=False nests the unshard/reshard override (via _fsdp2_unshard_context). +@pytest.mark.parametrize("writeback", [True, False]) +def test_persistent_materialization_only_skips_unshard_for_its_layer(cpu_dist_workers, writeback): + cpu_dist_workers.run(_forward_through_other_layers_while_one_is_materialized, writeback) + + +def _shared_group_override_restores_prior_attributes(rank, size): + mesh = init_device_mesh("cpu", (size,)) + block = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8)) + fully_shard(list(block), mesh=mesh) # both linears share one FSDPParamGroup + fully_shard(block, mesh=mesh) + group = fully_shard.state(block[0])._fsdp_param_group + prior_reshard = group.reshard # pre-existing instance override must survive + group.reshard = prior_reshard + with _disable_fsdp_unshard_reshard(block): + with _disable_fsdp_unshard_reshard(block): + pass + assert group.unshard() is None # outer no-op survives the inner exit + assert "unshard" in group.__dict__ + assert "unshard" not in group.__dict__ + assert group.__dict__["reshard"] is prior_reshard + # A layer that is not itself FSDP-wrapped was not gathered, so its FSDP submodules must unshard. + with _disable_fsdp_unshard_reshard(torch.nn.Sequential(*block)): + assert "unshard" not in group.__dict__ + + +def test_disable_fsdp_unshard_reshard_shared_group_and_nesting(cpu_dist_workers): + cpu_dist_workers.run(_shared_group_override_restores_prior_attributes) diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py index 931a4a3ee46..97def429a1e 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py +++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py @@ -361,7 +361,7 @@ def test_overfits_a_single_batch(self): input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size) optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=5e-3) - for _ in range(60): + for _ in range(30): # both accuracies pass 0.9 by step ~21 out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) optimizer.zero_grad() out.loss.backward() diff --git a/tests/unit/torch/utils/test_get_live_pool.py b/tests/unit/torch/utils/test_get_live_pool.py new file mode 100644 index 00000000000..2166d5ee560 --- /dev/null +++ b/tests/unit/torch/utils/test_get_live_pool.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``get_live_pool`` hands back a live pool and replaces a killed one.""" + +from types import SimpleNamespace + +from _test_utils.torch.distributed.utils import get_live_pool + + +def test_a_killed_pool_is_replaced_and_a_live_one_is_reused(): + pools = {} + made = [] + + def make(): + made.append(SimpleNamespace(broken=False)) + return made[-1] + + first = get_live_pool(pools, 2, make) + assert get_live_pool(pools, 2, make) is first + first.broken = True + second = get_live_pool(pools, 2, make) + assert second is not first and len(made) == 2