Refactor mdlm, bd3lm, ppl tracking, and minor improvements

[ZHZisZZ]

This pull request refactors and improves the BD3LM and MDLM trainer and sampler code, focusing on configuration management, loss computation, and sampling robustness. The changes introduce dataclass-based configuration, clarify and correct masking logic, and enhance per-sequence EOS handling in the sampler. Additionally, metric tracking is unified and modernized. The most important changes are grouped below:

Configuration and Trainer Refactoring:

• Introduced MDLMConfig and BD3LMConfig dataclasses for structured configuration of trainers, replacing multiple positional arguments and improving clarity and maintainability (dllm/core/trainers/mdlm.py, dllm/core/trainers/bd3lm.py). [1] [2]
• Updated trainer initialization to use the new config dataclasses and to unify metric tracking via a new OnEvaluateMetricsCallback instead of the previous EpochPPLMeter (dllm/core/trainers/mdlm.py, dllm/core/trainers/bd3lm.py). [1] [2]

Loss Computation and Masking Logic:

• Standardized masking logic by renaming variables (e.g., masked_indices to masked_mask, and token_cnt_per_seq to maskable_mask), and updated all related logic to use these clearer names. This change ensures consistency and correctness in how loss is computed and normalized (dllm/core/trainers/mdlm.py, dllm/core/trainers/bd3lm.py). [1] [2] [3] [4] [5]
• Improved degenerate case handling (when no tokens are masked) to ensure gradients remain valid, and updated metric tracking to use the new unified callback (dllm/core/trainers/bd3lm.py).
• Updated loss normalization to use the new config field loss_norm_type and fixed normalization logic for batch, sequence, and token modes (dllm/core/trainers/bd3lm.py).

Sampler Improvements:

• Enhanced block sampling in BD3LM sampler to pad prompt lengths to block size multiples, ensuring proper alignment and simplifying block scheduling (dllm/core/samplers/bd3lm.py). [1] [2]
• Added robust per-sequence EOS handling to allow early stopping of generation for sequences that have reached EOS, improving efficiency and correctness (dllm/core/samplers/bd3lm.py). [1] [2]
• Renamed internal functions for clarity (e.g., build_staircase_attention_mask to _prepare_for_sampling, diffusion_step_block to _diffusion_step_block) (dllm/core/samplers/bd3lm.py). [1] [2]

Documentation and API Consistency:

• Updated example SFT preprocessing commands in README.md to use the correct mapping function (dllm.utils.default_sft_map_fn) for consistency with code changes (README.md). [1] [2]

Other Code Quality Improvements:

• Removed unused imports and improved type annotations for better code clarity and maintainability (dllm/core/trainers/bd3lm.py). [1] [2]
• Renamed and refactored attention mask construction function for BD3LM training for better clarity (dllm/core/trainers/bd3lm.py). [1] [2] [3]

These changes collectively modernize the codebase, improve maintainability, and ensure more robust and correct training and sampling behavior.

[gemini-code-assist]

Summary of Changes

Hello @ZHZisZZ, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request primarily focuses on a significant refactoring of the training and sampling components, aiming to improve code structure, consistency, and maintainability. Key changes include standardizing trainer configurations using dataclasses, overhauling the metric logging system with torchmetrics for more robust and DDP-safe reporting, and enhancing the BD3LM sampler with better padding and sequence termination logic. These changes streamline development and ensure a more reliable and extensible framework.

Highlights

• Trainer Configuration Refactor: Introduced dedicated Config dataclasses (e.g., MDLMConfig, BD3LMConfig, DreamConfig, EditFlowConfig) within each trainer, inheriting from a base TrainingArguments class. This centralizes and standardizes trainer-specific arguments, improving clarity and maintainability across different training pipelines.
• Metric Logging System Overhaul: Replaced the custom EpochPPLMeter with a more robust and flexible BaseMetricsCallback and OnEvaluateMetricsCallback system, leveraging torchmetrics. This new system provides DDP-safe metric accumulation, per-split metric independence, and standardized logging for NLL and Perplexity.
• BD3LM Sampler Enhancements: Improved the bd3lm sampler by renaming internal utility functions for clarity, adding per-sequence End-of-Sequence (EOS) stopping logic, and refining prompt padding to align with block boundaries, leading to more efficient and controlled generation.
• Codebase Cleanup and Renaming: Numerous internal utility functions and variables were renamed (e.g., build_staircase_attention_mask to _prepare_for_sampling, default_mdlm_sft_map_fn to default_sft_map_fn, masked_indices to masked_mask) to enhance code readability and consistency. GPT2 A2D model support was also removed.
• New Dataset Utility and Dependency: Added a new utility (dllm/data/s1k.py) for loading and mapping the 's1K' dataset. The pyproject.toml was updated to include torchmetrics as a core dependency and a new rl optional dependency group.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[ZHZisZZ]

@copilot Please check for potential bugs and verify whether the changes are functionally equivalent before and after.

[Copilot]

@ZHZisZZ I've opened a new pull request, #69, to work on those changes. Once the pull request is ready, I'll request review from you.

[gemini-code-assist]

Code Review

This pull request introduces a significant and high-quality refactoring across the codebase. Key improvements include:

• A shift to structured dataclass-based configurations for trainers, enhancing maintainability.
• The implementation of a new, robust metrics system using torchmetrics, replacing the older manual approach.
• Cleaner and more efficient loss computation logic in the trainers.
• Numerous code cleanups, such as function renaming, deduplication, and making example scripts self-contained.
• Correctness improvements in the block diffusion sampler for batch generation.

The overall changes greatly improve the quality, clarity, and maintainability of the code. I have a few suggestions for further improvement, including a potential correctness issue regarding a removed chat template and a couple of minor style fixes.

[gemini-code-assist]

This refactoring combines the logic for A2DQwen2LMHeadModel and A2DQwen3LMHeadModel, which is great for deduplication. However, the very long and specific chat_template that was previously set for A2DQwen3LMHeadModel has been removed. This template is crucial for the model's chat and tool-use functionality. Its removal will likely lead to incorrect behavior. If this template is not being set by some other mechanism (e.g., from the tokenizer's config on the Hub), it should be restored, perhaps conditionally for A2DQwen3LMHeadModel.

[gemini-code-assist]

This commented-out block is a valuable optimization. By filling new blocks with pad_id for sequences that are already done, you can avoid wasteful computation during the diffusion steps. The current implementation processes these completed sequences unnecessarily. I recommend re-enabling this logic to improve efficiency in batch generation.

Suggested change

	# if done.any():
	# new_block = torch.where(
	# done.unsqueeze(1),
	# torch.full_like(new_block, pad_id),
	# new_block,
	# )
	if done.any():
	new_block = torch.where(
	done.unsqueeze(1),
	torch.full_like(new_block, pad_id),
	new_block,
	)

[gemini-code-assist]

This file is missing a final newline character. It's a common Python convention (recommended by PEP 8) to end files with a single newline. This improves consistency and can prevent issues with some tools and file concatenations.

[gemini-code-assist]

This file is missing a final newline character. It's a common Python convention (recommended by PEP 8) to end files with a single newline. This improves consistency and can prevent issues with some tools and file concatenations.