Implementation:NVIDIA NeMo Aligner MegatronSDXL DRaFTP Model
| Knowledge Sources | |
|---|---|
| Domains | Multimodal, Image Generation, Diffusion Models, Alignment |
| Last Updated | 2026-02-08 00:00 GMT |
Overview
MegatronSDXLDRaFTPModel implements the DRaFT+ alignment method for Stable Diffusion XL (SDXL) models, extending MegatronDiffusionEngine with truncated backpropagation, KL regularization, FSDP support, and PEFT (LoRA) compatibility.
Description
The MegatronSDXLDRaFTPModel class extends MegatronDiffusionEngine and implements SupervisedInterface. It is specifically designed for SDXL models, which use a different architecture and API compared to standard SD models.
Key architectural differences from MegatronSDDRaFTPModel:
- Conditional model: Uses SDXL's conditioner system with unconditional conditioning and zero-forced embeddings for text/captions, rather than the simpler CLIP encoding of standard SD.
- PEFT support: Supports both full fine-tuning and PEFT (LoRA). When peft_scheme is "none" (full fine-tuning), a frozen copy of the base DiffusionEngine is created. When using LoRA, the base model is accessed via adapter_control context manager which temporarily disables adapters.
- FSDP sharding: Overrides configure_sharded_model to handle FSDP with frozen submodules (VAE, CLIP embedders) properly excluded from sharding via ignored_states.
- VAE optimization: Deletes the VAE encoder after initialization since only the decoder is needed during training/inference.
- Sampler: Uses get_sampler_config from SDXL helpers instead of initializing a DDIM sampler directly.
- Size conditioning: The append_sdxl_size_keys method adds SDXL-specific batch keys (original_size_as_tuple, target_size_as_tuple, crop_coords_top_left).
The generate method follows the same truncated backpropagation pattern as the SD variant, but uses the SDXL sampler API (sampler.prepare_sampling_loop, sampler.sampler_step, sampler.get_gamma) instead of DDIM steps.
The annealed_guidance method similarly interpolates between base and fine-tuned denoiser functions at each step, using the SDXL sampling loop with Euler step integration.
Usage
Instantiate this model for DRaFT+ training on SDXL models. Assign a reward model to ptl_model.reward_model before training. Supports both full fine-tuning and LoRA via PEFT configuration.
Code Reference
Source Location
- Repository: NVIDIA_NeMo_Aligner
- File: nemo_aligner/models/mm/stable_diffusion/megatron_sdxl_draftp_model.py
- Lines: 75-559
Signature
class MegatronSDXLDRaFTPModel(MegatronDiffusionEngine, SupervisedInterface):
def __init__(self, cfg, trainer):
def configure_sharded_model(self):
def append_sdxl_size_keys(self, prompts):
def generate(self, batch, x_T):
def annealed_guidance(self, batch, x_T, weighing_fn=None):
def get_forward_output_and_loss_func(self, validation_step=False):
def get_loss_and_metrics(self, batch, forward_only=False):
Import
from nemo_aligner.models.mm.stable_diffusion.megatron_sdxl_draftp_model import MegatronSDXLDRaFTPModel
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| cfg | DictConfig | Yes | Model configuration with sampling.base (height, width), truncation_steps, kl_coeff, peft.peft_scheme, first_stage_config, micro_batch_size, global_batch_size, precision |
| trainer | Trainer | Yes | PyTorch Lightning Trainer instance |
| batch | list[str] | Yes | List of text prompts for image generation |
| x_T | Tensor | Yes | Initial noise latents of shape [B, C, H//f, W//f] |
Outputs
| Name | Type | Description |
|---|---|---|
| loss_value | float | Negative mean reward plus KL penalty |
| metrics | dict | Dictionary with "loss" and "kl_penalty" values; during validation also includes "images_and_captions" |
| image | Tensor | Generated images scaled to [0, 255] of shape [B, C, H, W] (from generate) |
| t_eps_draft_p | Tensor | Stacked noise predictions from fine-tuned model over truncated steps |
| t_eps_init | Tensor | Stacked noise predictions from base model over truncated steps |
Usage Examples
from nemo_aligner.models.mm.stable_diffusion.megatron_sdxl_draftp_model import MegatronSDXLDRaFTPModel
from nemo_aligner.models.mm.stable_diffusion.image_text_rms import get_reward_model
# Instantiate SDXL DRaFT+ model
ptl_model = MegatronSDXLDRaFTPModel(cfg.model, trainer).to(torch.cuda.current_device())
init_peft(ptl_model, cfg.model) # Initialize LoRA if configured
# Attach reward model
reward_model = get_reward_model(cfg.rm, mbs=cfg.model.micro_batch_size, gbs=cfg.model.global_batch_size)
ptl_model.reward_model = reward_model
# Training step
loss_value, metrics = ptl_model.get_loss_and_metrics(batch, forward_only=False)
# Annealed inference with linear weighting
images = ptl_model.annealed_guidance(prompts, latents, weighing_fn=lambda s1, s2, i, t: i / t)