From efdd676d5678be7b5112443dc44cd5cde0daa461 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 13 Feb 2026 21:18:02 +0800 Subject: [PATCH] [diffusion] refactor: merge redundant default_dtype and param_dtype parameters in FSDP loader (#18789) --- .../loader/component_loaders/bridge_loader.py | 3 +-- .../component_loaders/transformer_loader.py | 1 - .../multimodal_gen/runtime/loader/fsdp_load.py | 16 +++++++++++----- 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py index da084ec19..29f216fec 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py @@ -88,8 +88,7 @@ class BridgeLoader(ComponentLoader): cpu_offload=server_args.dit_cpu_offload, pin_cpu_memory=server_args.pin_cpu_memory, fsdp_inference=server_args.use_fsdp_inference, - default_dtype=default_dtype, - param_dtype=torch.bfloat16, + param_dtype=default_dtype, reduce_dtype=torch.float32, output_dtype=None, strict=False, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index ee7b33424..279b0b4ea 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py @@ -105,7 +105,6 @@ class TransformerLoader(ComponentLoader): pin_cpu_memory=server_args.pin_cpu_memory, fsdp_inference=server_args.use_fsdp_inference, # TODO(will): make these configurable - default_dtype=default_dtype, param_dtype=torch.bfloat16, reduce_dtype=torch.float32, output_dtype=None, diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 2c723ab27..af10e8e31 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -54,7 +54,6 @@ def maybe_load_fsdp_model( device: torch.device, hsdp_replicate_dim: int, hsdp_shard_dim: int, - default_dtype: torch.dtype, param_dtype: torch.dtype, reduce_dtype: torch.dtype, cpu_offload: bool = False, @@ -63,8 +62,15 @@ def maybe_load_fsdp_model( pin_cpu_memory: bool = True, strict: bool = True, ) -> torch.nn.Module: - """ - Load the model with FSDP if is training, else load the model without FSDP. + """Load a model with optional FSDP (Fully Sharded Data Parallel) support. + + Args: + param_dtype: Data type for model parameters, also used for: + - Model initialization context (set_default_torch_dtype) + - FSDP mixed precision policy + - Weight loading and casting + reduce_dtype: Data type for gradient reduction in FSDP mixed precision. + strict: If True, enforce strict state dict loading (all keys must match). """ # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are # manually casting the inputs to the model @@ -79,7 +85,7 @@ def maybe_load_fsdp_model( mp_policy=mp_policy, ) - with set_default_torch_dtype(default_dtype), torch.device("meta"): + with set_default_torch_dtype(param_dtype), torch.device("meta"): model = model_cls(**init_params) # Check if we should use FSDP @@ -120,7 +126,7 @@ def maybe_load_fsdp_model( model, weight_iterator, device, - default_dtype, + param_dtype, strict=strict, cpu_offload=cpu_offload, param_names_mapping=param_names_mapping_fn,