[diffusion] refactor: add arg to control the precision of dit (#17751)

python/sglang/multimodal_gen/docs/cli.md

@@ -23,6 +23,7 @@ The SGLang-diffusion CLI provides a quick way to access the inference pipeline f
 - `--ring-degree {RING_DEGREE}`: The degree of ring attention-style SP in USP
 - `--attention-backend {BACKEND}`: Attention backend to use. For SGLang-native pipelines use `fa`, `torch_sdpa`, `sage_attn`, etc. For diffusers pipelines use diffusers backend names like `flash`, `_flash_3_hub`, `sage`, `xformers`.
 - `--cache-dit-config {PATH}`: Path to a Cache-DiT YAML/JSON config (diffusers backend only)
+- `--dit-precision {DTYPE}`: Precision for the DiT model (currently supports fp32, fp16, and bf16).
 
 
 ### Sampling Parameters

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -6,7 +6,6 @@
 # Copyright 2024 The TorchTune Authors.
 # Copyright 2025 The sglang-diffusion Authors.
 
-import contextlib
 from collections.abc import Callable, Generator
 from itertools import chain
 from typing import Any
@@ -26,6 +25,7 @@ from torch.nn.modules.module import _IncompatibleKeys
 from sglang.multimodal_gen.runtime.loader.utils import (
     get_param_names_mapping,
     hf_to_custom_state_dict,
+    set_default_torch_dtype,
 )
 from sglang.multimodal_gen.runtime.loader.weight_utils import (
     safetensors_weights_iterator,
@@ -46,34 +46,6 @@ def _make_param_like(
     return new_param
 
 
-# TODO(PY): move this to utils elsewhere
-@contextlib.contextmanager
-def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
-    """
-    Context manager to set torch's default dtype.
-
-    Args:
-        dtype (torch.dtype): The desired default dtype inside the context manager.
-
-    Returns:
-        ContextManager: context manager for setting default dtype.
-
-    Example:
-        >>> with set_default_dtype(torch.bfloat16):
-        >>>     x = torch.tensor([1, 2, 3])
-        >>>     x.dtype
-        torch.bfloat16
-
-
-    """
-    old_dtype = torch.get_default_dtype()
-    torch.set_default_dtype(dtype)
-    try:
-        yield
-    finally:
-        torch.set_default_dtype(old_dtype)
-
-
 # TODO(PY): add compile option
 def maybe_load_fsdp_model(
     model_cls: type[nn.Module],
@@ -107,7 +79,7 @@ def maybe_load_fsdp_model(
         mp_policy=mp_policy,
     )
 
-    with set_default_dtype(default_dtype), torch.device("meta"):
+    with set_default_torch_dtype(default_dtype), torch.device("meta"):
         model = model_cls(**init_params)
 
     # Check if we should use FSDP

python/sglang/multimodal_gen/runtime/loader/utils.py

@@ -20,8 +20,10 @@ def set_default_torch_dtype(dtype: torch.dtype):
     """Sets the default torch dtype to the given dtype."""
     old_dtype = torch.get_default_dtype()
     torch.set_default_dtype(dtype)
-    yield
-    torch.set_default_dtype(old_dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(old_dtype)
 
 
 def get_param_names_mapping(

python/sglang/multimodal_gen/runtime/pipelines/comfyui_zimage_pipeline.py

@@ -15,10 +15,12 @@ from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.loader.fsdp_load import (
     load_model_from_full_model_state_dict,
-    set_default_dtype,
     shard_model,
 )
-from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
+from sglang.multimodal_gen.runtime.loader.utils import (
+    get_param_names_mapping,
+    set_default_torch_dtype,
+)
 from sglang.multimodal_gen.runtime.loader.weight_utils import (
     safetensors_weights_iterator,
 )
@@ -299,7 +301,7 @@ class ComfyUIZImagePipeline(LoRAPipeline, ComposedPipelineBase):
                 mp_policy=mp_policy,
             )
 
-            with set_default_dtype(default_dtype), torch.device("meta"):
+            with set_default_torch_dtype(default_dtype), torch.device("meta"):
                 model = model_cls(**{"config": dit_config, "hf_config": hf_config})
 
             # Check if we should use FSDP