Skip to content

Restore enforce_marked_param_dtypes after model materialization - #60

Merged
yueming-yuan merged 1 commit into
miles-mainfrom
restore-enforce-marked-param-dtypes
Jun 29, 2026
Merged

Restore enforce_marked_param_dtypes after model materialization#60
yueming-yuan merged 1 commit into
miles-mainfrom
restore-enforce-marked-param-dtypes

Conversation

@guapisolo

@guapisolo guapisolo commented Jun 29, 2026

Copy link
Copy Markdown

What

Restores the call to enforce_marked_param_dtypes(model) in get_model (megatron/training/training.py), right after correct_amax_history_if_needed(model) and before the DDP wrap.

# Miles allows selected parameters to opt out of Float16Module's global
# bf16/fp16 cast. Restore those dtypes immediately after model materialization.
from miles.backends.megatron_utils.fp32_param_utils import enforce_marked_param_dtypes

enforce_marked_param_dtypes(model)

Why

PR #28 (DeepSeek V4 RL support) removed this block. Miles allows selected parameters to opt out of Float16Module's global bf16/fp16 cast; this call restores those dtypes immediately after model materialization. Without it, marked fp32 params get silently cast, so the block is added back.

This change is an exact revert of the removal in PR #28 — the resulting file is byte-identical to the pre-#28 state.

Acknowledgements

Thanks to @zyzshishui for finding this bug. 🙏

🤖 Generated with Claude Code

PR #28 (DeepSeek V4 RL support) removed the call to
enforce_marked_param_dtypes in get_model. Miles allows selected
parameters to opt out of Float16Module's global bf16/fp16 cast, and
this call restores those dtypes immediately after model
materialization, so it is added back.

Thanks to @zyzshishui for finding this bug.

Co-authored-by: zyzshishui <82826991+zyzshishui@users.noreply.github.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@guapisolo
guapisolo force-pushed the restore-enforce-marked-param-dtypes branch from a7342dc to 6cc7765 Compare June 29, 2026 08:27
@yueming-yuan
yueming-yuan merged commit 79fc089 into miles-main Jun 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants