fix(diffusion): apply compute_dtype via autocast when FSDP2 skips (#3589)

* fix(diffusion): apply compute_dtype via autocast when FSDP2 skips sharding

Split model.torch_dtype/model.compute_dtype relies on FSDP2's
MixedPrecisionPolicy to cast parameters to the compute dtype. On a
single-rank world or mesh, FSDP2Manager.parallelize skips sharding, so
that cast never happens: parameters stay in torch_dtype while the
flow-matching step still casts its inputs to compute_dtype. The forward
pass then mixes dtypes and fails with

  RuntimeError: Input type (c10::BFloat16) and bias type (float) should
  be the same

on the transformer's fp32 patch embedding. The same config works on 2+
GPUs, and five shipped diffusion example configs use this split
(torch_dtype: float32 with compute_dtype: bfloat16).

Wrap the forward pass in torch.autocast on that path, following the
autocast approach the retrieval and dLLM recipes already use for
FSDP2Config.autocast_dtype. Resident parameters and their gradients stay
in torch_dtype, so optimizer state is unchanged. Numerics are close to
but not identical to the sharded path, since autocast keeps reductions
such as normalization in fp32 rather than casting every parameter.

Extract the skip condition into fsdp2_sharding_enabled() so the recipe
checks the same predicate FSDP2Manager acts on instead of duplicating it.

LLM/VLM recipes drop mp_policy at world size 1 the same way, but degrade
silently to fp32 compute rather than failing; that is left for a separate
change.

Fixes #3578

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

* test(diffusion): stub _autocast_dtype in the metrics recipe test

test_run_train_validation_loop_uses_hot_path_and_logs_perf_metrics builds the
recipe with object.__new__ and assigns each attribute the training loop reads,
bypassing setup(). The loop now also enters _autocast_context(), which reads
self._autocast_dtype, so the test failed with AttributeError on both
L0_Unit_Tests_CPU and L0_Unit_Tests_GPU_1_of_2.

Set it to None alongside the existing transformer_engine_fp8=False stub, which
is the direct analogue: both disable a per-forward wrapping context. The recipe
class intentionally has no class-level attribute defaults -- setup() owns every
attribute -- so the stub belongs in the test rather than in the recipe.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

---------

Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>