Skip to content

[diffusion] UX: aggregate expected dtype-cast logs during weight loading #21552

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mickqian merged 2 commits into from
Mar 28, 2026
Merged

[diffusion] UX: aggregate expected dtype-cast logs during weight loading #21552

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4221261
[diffusion] UX: aggregate expected dtype-cast logs during weight loading
mickqian Mar 27, 2026
6e3668d
relax
mickqian Mar 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 74 additions & 20 deletions python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
# Copyright 2024 The TorchTune Authors.
# Copyright 2025 The sglang-diffusion Authors.

from collections import Counter, defaultdict
from collections.abc import Callable, Generator
from itertools import chain
from typing import Any
Expand Down Expand Up @@ -40,6 +41,28 @@

logger = init_logger(__name__)

_QUANTIZED_DTYPES = (
torch.uint8,
torch.float8_e4m3fn,
torch.float8_e5m2,
torch.int8,
)
_DTYPE_MISMATCH_EXAMPLE_LIMIT = 3


def _format_dtype_mismatch_summary(
mismatch_counts: Counter[tuple[torch.dtype, torch.dtype]],
mismatch_examples: dict[tuple[torch.dtype, torch.dtype], list[str]],
) -> str:
parts: list[str] = []
for (checkpoint_dtype, target_dtype), count in mismatch_counts.items():

@gemini-code-assist gemini-code-assist Bot Mar 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

For deterministic logging, it's good practice to sort the items from the mismatch_counts dictionary. This ensures the output summary is always in the same order, which can be helpful for debugging and testing. You can sort by the string representation of the dtypes, as torch.dtype objects are not directly comparable.

Suggested change
for (checkpoint_dtype, target_dtype), count in mismatch_counts.items():
for (checkpoint_dtype, target_dtype), count in sorted(mismatch_counts.items(), key=lambda item: (str(item[0][0]), str(item[0][1]))):

examples = mismatch_examples[(checkpoint_dtype, target_dtype)]
part = f"{checkpoint_dtype}->{target_dtype} x{count}"
if examples:
part += f" (e.g. {', '.join(examples)})"
parts.append(part)
return "; ".join(parts)


def _make_param_like(
actual_param: torch.nn.Parameter, tensor: torch.Tensor
Expand Down Expand Up @@ -272,6 +295,18 @@ def load_model_from_full_model_state_dict(

sharded_sd = {}
skipped_checkpoint_keys: list[str] = []
non_quantized_dtype_mismatch_counts: Counter[tuple[torch.dtype, torch.dtype]] = (
Counter()
)
non_quantized_dtype_mismatch_examples: dict[
tuple[torch.dtype, torch.dtype], list[str]
] = defaultdict(list)
quantized_dtype_mismatch_counts: Counter[tuple[torch.dtype, torch.dtype]] = (
Counter()
)
quantized_dtype_mismatch_examples: dict[
tuple[torch.dtype, torch.dtype], list[str]
] = defaultdict(list)

# shard from loaded state_dict, custom_param_sd -> sharded_sd
for target_param_name in sorted_param_names:
Expand All @@ -296,32 +331,29 @@ def load_model_from_full_model_state_dict(
else:
target_dtype = meta_sharded_param.dtype

_QUANTIZED_DTYPES = (
torch.uint8,
torch.float8_e4m3fn,
torch.float8_e5m2,
torch.int8,
)
if full_tensor.dtype != target_dtype:
mismatch_key = (full_tensor.dtype, target_dtype)
if (
full_tensor.dtype in _QUANTIZED_DTYPES
or target_dtype in _QUANTIZED_DTYPES
):
logger.warning(
"Dtype mismatch for quantized parameter %s: "
"checkpoint has %s, model expects %s",
target_param_name,
full_tensor.dtype,
target_dtype,
)
quantized_dtype_mismatch_counts[mismatch_key] += 1
if (
len(quantized_dtype_mismatch_examples[mismatch_key])
< _DTYPE_MISMATCH_EXAMPLE_LIMIT
):
quantized_dtype_mismatch_examples[mismatch_key].append(
target_param_name
)
else:
logger.warning(
"Dtype mismatch for %s: checkpoint has %s, model expects %s. "
"Casting checkpoint tensor to the target dtype during load.",
target_param_name,
full_tensor.dtype,
target_dtype,
)
non_quantized_dtype_mismatch_counts[mismatch_key] += 1
if (
len(non_quantized_dtype_mismatch_examples[mismatch_key])
< _DTYPE_MISMATCH_EXAMPLE_LIMIT
):
non_quantized_dtype_mismatch_examples[mismatch_key].append(
target_param_name
)

if not hasattr(meta_sharded_param, "device_mesh"):
full_tensor = full_tensor.to(device=device, dtype=target_dtype)
Expand Down Expand Up @@ -378,6 +410,28 @@ def load_model_from_full_model_state_dict(

model.reverse_param_names_mapping = reverse_param_names_mapping

if non_quantized_dtype_mismatch_counts:
logger.debug(
"Casting checkpoint tensors to target dtype during load: %s",
_format_dtype_mismatch_summary(
non_quantized_dtype_mismatch_counts,
non_quantized_dtype_mismatch_examples,
),
main_process_only=True,
local_main_process_only=True,
)

if quantized_dtype_mismatch_counts:
logger.warning(
"Dtype mismatches detected for quantized parameters during load: %s",
_format_dtype_mismatch_summary(
quantized_dtype_mismatch_counts,
quantized_dtype_mismatch_examples,
),
main_process_only=True,
local_main_process_only=True,
)

if skipped_checkpoint_keys:
logger.warning(
"Checkpoint keys not loaded (no matching model parameter) %s",
Expand Down
Loading