Add multi-precision training support to FSDP script

[aagallo]

Description

This PR adds comprehensive precision parameter support to the FSDP training script, enabling users to configure training with multiple precision formats (FP32, FP16, FP8, MXFP8, NVFP4) via command-line argument. The implementation includes automatic configuration of appropriate dtypes and format-specific recipes for each precision type.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Added precision() type validator function supporting fp32, fp16, fp8, mxfp8, and nvfp4 formats
• Added --precision command-line argument to parse_fsdp_args() with default value "fp8"
• Implemented match statement in train() function to configure precision-based training parameters
• Configured format-specific recipes for each precision type:
  • FP32/FP16: Uses standard PyTorch dtypes with FP8 disabled
  • FP8: Uses DelayedScaling recipe with HYBRID format
  • MXFP8: Uses MXFP8BlockScaling recipe with E4M3 format
  • NVFP4: Uses NVFP4BlockScaling recipe with bfloat16 dtype
• Set appropriate no_fp8 flags based on precision selection
• Updated layer_kwargs["params_dtype"] to use precision-determined dtype
• Imported required recipe classes: MXFP8BlockScaling and NVFP4BlockScaling

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

Please reach out to Santosh Bhavani ([email protected]) for additional context on the work

[greptile-apps]

Greptile Overview

Greptile Summary

This PR extends the PyTorch FSDP example (examples/pytorch/fsdp/fsdp.py) with a --precision preset flag that configures parameter dtype, FP8 enablement, and the appropriate Transformer Engine recipe for FP8/MXFP8/NVFP4 modes.

The script now tracks whether --dtype / --no-fp8 were explicitly provided, applies documented precedence rules, and then propagates the resolved dtype/recipe through TE layer construction (params_dtype), FSDP mixed precision (MixedPrecision.param_dtype), input tensor creation, and the te.autocast() context.

Confidence Score: 5/5

• This PR is safe to merge with minimal risk.
• Reviewed the full diff in the only changed file and traced the new recipe classes in transformer_engine.common.recipe. The updated precedence logic yields consistent dtype/no_fp8/recipe selection and ensures recipe is defined for all branches. No definite runtime errors or broken behavior introduced in the latest commit were found beyond issues already covered in prior review threads.
• examples/pytorch/fsdp/fsdp.py

Important Files Changed

Filename	Overview
examples/pytorch/fsdp/fsdp.py	Adds --precision presets with explicit-flag precedence tracking, selects dtype/FP8 recipe accordingly, and propagates dtype into TE layer params, FSDP MixedPrecision, inputs, and te.autocast recipe.

Sequence Diagram

sequenceDiagram
    participant U as User/CLI
    participant P as parse_fsdp_args()
    participant T as train(opts)
    participant TE as te.autocast
    participant F as FSDP

    U->>P: Run examples/pytorch/fsdp/fsdp.py with flags
    P-->>T: opts incl. precision/dtype/no_fp8 + explicit-set markers
    T->>T: Compute dtype, no_fp8, recipe (preset + overrides)
    T->>F: Wrap model with MixedPrecision(param_dtype=dtype)

    loop Each iteration
        T->>T: Create input tensor x (dtype)
        T->>TE: Enter autocast (enabled = not no_fp8)
        TE-->>T: Forward executes with recipe
        T->>T: Backward + optimizer step
    end

Loading

[greptile-apps]

_{1 file reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Incorrect fp16 dtype
In the case "fp16" branch, the code sets dtype = torch.bfloat16. That contradicts the meaning of fp16 and also diverges from the existing --dtype parsing which supports torch.float16. If a user runs with --precision fp16 expecting fp16 parameters/inputs, they’ll silently get bf16 instead.

[greptile-apps]

Conflicting CLI flags
--precision and --dtype/--no-fp8 now overlap: train() overrides dtype and no_fp8 based on --precision, but --dtype/--no-fp8 are still accepted and used as defaults. As written, --precision fp8 will force no_fp8=False even if the user explicitly passed --no-fp8, and --precision fp32 will ignore an explicit --dtype fp16/bf16. This makes the CLI behavior surprising and hard to reason about; either make --precision derive defaults only when the user didn’t specify --dtype/--no-fp8, or document/enforce precedence (e.g., error on incompatible combinations).

[greptile-apps]

Additional Comments (1)

examples/pytorch/fsdp/fsdp.py
FSDP mixed_precision mismatch
layer_kwargs["params_dtype"] and the input tensor dtype are overridden by --precision, but FSDP is still configured with mixed_precision=MixedPrecision(param_dtype=opts.dtype, ...) (and opts.dtype no longer matches the model param dtype when --precision is used). This will cause inconsistent param casting/communication behavior under FSDP for e.g. --precision fp32 (params are fp32 but FSDP thinks they’re bf16) and --precision fp16 (currently sets dtype=torch.bfloat16). FSDP param_dtype should be driven by the same dtype selected in the precision switch, or the precision switch should not override param dtype when FSDP mixed precision is enabled.

[greptile-apps]

_{1 file reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

recipe variable not initialized in default case. If precision validator allows an unexpected value, this will cause UnboundLocalError at line 403 when passed to te.autocast().

Suggested change

	case _:
	dtype = torch.float16
	precision_format = Format.HYBRID
	recipe = DelayedScaling(
	fp8_format=precision_format, amax_history_len=32, amax_compute_algo="max"
	)
	no_fp8 = opts.no_fp8
	case _:
	dtype = torch.float16
	precision_format = Format.HYBRID
	recipe = DelayedScaling(
	fp8_format=precision_format, amax_history_len=32, amax_compute_algo="max"
	)
	no_fp8 = False

[greptile-apps]

Parsing sys.argv directly is fragile - will break if args are passed via config file, environment variables, or if arg uses = syntax (--dtype=fp16). Use parser.parse_known_args() or track which args were explicitly set via custom action class.

Example with custom action:

class StoreExplicitAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, values)
        setattr(namespace, f'{self.dest}_explicitly_set', True)

[greptile-apps]

Additional Comments (1)

examples/pytorch/fsdp/fsdp.py
FSDP MixedPrecision uses opts.dtype instead of local dtype variable. When --precision is used without explicit --dtype, FSDP will use the default bfloat16 instead of the precision-determined dtype (e.g., float16 for fp8/mxfp8 or float32 for fp32).

        mixed_precision=MixedPrecision(
            param_dtype=dtype,
            reduce_dtype=torch.float32,
        ),

[greptile-apps]

_{1 file reviewed, 5 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Detection of explicit --no-fp8 is incorrect - opts.no_fp8 != False is True when the flag is absent (defaults to False). User running with defaults will trigger else branch instead of precision preset logic.

Suggested change

	no_fp8_explicitly_set = opts.no_fp8 != False
	no_fp8_explicitly_set = "--no-fp8" in sys.argv

[greptile-apps]

Missing action=StoreExplicitAction to track explicit --no-fp8 usage - required for precedence logic at line 250

Suggested change

	"--no-fp8",
	action=StoreExplicitAction,

[greptile-apps]

Logging FP8 recipe when no_fp8=True and recipe=None (line 322) will cause AttributeError on type(recipe).__name__

[greptile-apps]

Default changed from implicit FP8 disabled (via --dtype bf16 + --no-fp8 False) to explicit --precision fp8 enabled - existing scripts relying on default behavior will now use FP8 unexpectedly

[greptile-apps]

Additional Comments (1)

examples/pytorch/fsdp/fsdp.py
param_dtype set to opts.dtype instead of local dtype variable - ignores precision preset

            param_dtype=dtype,

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

examples/pytorch/fsdp/fsdp.py
Broken --dtype parsing

--dtype uses a custom action=StoreExplicitAction, but StoreExplicitAction.__call__() just assigns values directly and never applies the type=torch_dtype conversion. With a custom Action, argparse won’t run the type converter for you, so passing --dtype bf16 will make opts.dtype a raw string (e.g., 'bf16') and later calls like torch.rand(..., dtype=dtype) / MixedPrecision(param_dtype=dtype) will fail.

One fix is to perform conversion inside the action:

    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, torch_dtype(values))
        setattr(namespace, f"{self.dest}_explicitly_set", True)

(or alternatively, implement a custom action that subclasses _StoreAction and preserves normal type handling).

[greptile-apps]

_{1 file reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

CRITICAL BUG: Incorrect recipe selection when --dtype or --no-fp8 is explicitly set

When users explicitly set --dtype but also specify --precision mxfp8 or --precision nvfp4, the code ignores the precision preset and falls back to DelayedScaling recipe instead of the expected MXFP8BlockScaling or NVFP4BlockScaling recipe.

Reproduction:

# User expects MXFP8BlockScaling recipe
./fsdp.py --precision mxfp8 --dtype bf16
# But gets DelayedScaling recipe instead!

Root cause:
Lines 340-344 unconditionally use DelayedScaling when not no_fp8 is true, completely ignoring the opts.precision value.

Expected behavior:
The code should respect the precision preset even when dtype is explicitly set. The logic should check opts.precision and select the appropriate recipe (MXFP8BlockScaling for "mxfp8", NVFP4BlockScaling for "nvfp4", etc.).

Suggested fix:

# Set up recipe based on precision preset
if not no_fp8:
    match opts.precision:
        case "mxfp8":
            recipe = MXFP8BlockScaling(fp8_format=Format.E4M3)
        case "nvfp4":
            recipe = NVFP4BlockScaling()
        case _:
            precision_format = Format.HYBRID
            recipe = DelayedScaling(
                fp8_format=precision_format, amax_history_len=32, amax_compute_algo="max"
            )
else:
    recipe = None

This is a merge-blocking bug as it breaks the advertised functionality of the --precision flag.

[greptile-apps]

Precision preset ignored

In the else: branch (when --dtype and/or --no-fp8 were explicitly set), the code unconditionally falls back to DelayedScaling whenever FP8 is enabled:

if not no_fp8:
    recipe = DelayedScaling(...)

This ignores opts.precision entirely, so --precision mxfp8 or --precision nvfp4 will silently use DelayedScaling if the user also sets --dtype/--no-fp8 (e.g. --precision mxfp8 --dtype bf16). That breaks the advertised presets and applies the wrong quantization recipe.

Fix: in the explicit-flags path, either (a) still select recipe based on opts.precision when FP8 is enabled, or (b) explicitly error/warn and force opts.precision back to fp8 if you’re going to always use DelayedScaling.

[greptile-apps]

_{1 file reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Unreachable default preset

In get_precision_preset(), the case _: branch defaults to FP8 behavior, but precision() already rejects any value outside {fp32, fp16, fp8, mxfp8, nvfp4}. As a result this branch should be dead code in normal execution, and keeping it can mask future bugs (e.g., if validation is accidentally bypassed). Consider replacing it with raise AssertionError/ValueError or removing the branch entirely so invalid values fail loudly.

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

[greptile-apps]

Redundant recipe recreation

When dtype_explicitly_set is true and FP8 remains enabled, the code overwrites recipe with get_recipe_for_precision(opts.precision), even though recipe was already set from the selected preset earlier. This creates a second recipe instance (and can diverge if preset recipe configuration changes). If the intent is “dtype override shouldn’t affect recipe”, you can keep the existing recipe rather than re-instantiating it.

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Missing default return

get_precision_preset() has no case _: / explicit raise, so it can return None and train() will crash on unpacking (preset_dtype, preset_no_fp8, preset_recipe = ...) if precision_value is ever outside the handled set (e.g., if validation is bypassed or a new preset is added but this match isn’t updated). Consider adding an explicit case _: that raises ValueError/AssertionError so invalid values fail loudly instead of a TypeError later.

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}