r"""Model-agnostic parts of Parallel Decoding Distillation (PDD). PDD (arXiv 2607.26004) turns a pre-trained flow model into a *parallel decoder*: the sampling interval is discretized into `N` intervals grouped into blocks of size `L`, and one network evaluation predicts the **mean velocity of every interval of the next block** instead of the single instantaneous velocity. Generation then advances `L` intervals per evaluation, i.e. `NFE = N / L`. Everything here works on plain `nn.Linear`s and tensors, with no reference to any particular transformer: the math of the time grid and the head plans, the [`PDDParallelHead`] / [`PDDLoRALinear`] modules, the teacher switch, and the checkpoint resolution. Porting a model supplies only the glue that knows where the final linear layers live and how a forward is called — `videox_fun/models/minimax_h3_pdd.py` is the MiniMax-H3 reference: an attach step that swaps the final heads for [`PDDParallelHead`]s, a teacher mean-velocity estimate in the model's calling convention, a step callback in the pipeline's callback protocol, and a `load_pdd_lora` that ties them together. """ import contextlib import json import os from typing import Optional, Sequence import torch import torch.nn as nn import torch.nn.functional as F def shifted_sigma(shift: float, sigma: torch.Tensor) -> torch.Tensor: r"""The exponential sigma shift of a rectified-flow schedule, `sigma' = s*sigma / (1 + (s-1)*sigma)`.""" return shift * sigma / (1 + (shift - 1) * sigma) def pdd_time_grid(shift: float, num_steps: int) -> torch.Tensor: r""" The PDD time discretization `0 = t_0 < ... < t_N = 1` of a rectified-flow schedule with an exponential sigma shift. The paper's shift reparameterization (eq. 16), `t_n = shift_s(n/N)` with `shift_s(t) = (t/s) / (1 + (1/s - 1) t)`, is algebraically the same grid as `t = 1 - sigma'` over a uniform sigma grid — the convention of MiniMax-H3's scheduler, where `t = 1` is clean. A consequence worth relying on when the model's scheduler is a plain Euler rectified-flow one: the block boundaries of this grid, taken every `L` indices, are exactly the grid `set_timesteps(N / L + 1)` builds, so PDD generation reuses such a released scheduler unchanged. Args: shift (`float`): The exponential shift of the schedule (`12.0` video / `3.0` audio for MiniMax-H3; `1.0` is a uniform grid). num_steps (`int`): The grid size `N`. Returns: `torch.Tensor` of shape `(num_steps + 1,)`, float64: the grid, ascending from `0` to `1`. """ sigma = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64) return 1.0 - shifted_sigma(shift, sigma) def pdd_training_plan(step_sizes: torch.Tensor, start: int, targets: Sequence[int], advance: int) -> torch.Tensor: r""" Every direction one PDD training step needs, from a single backbone evaluation. Args: step_sizes (`torch.Tensor` of shape `(N,)`): The grid step sizes `h_l = t_{l+1} - t_l`. start (`int`): The block start `n`, i.e. the index the state is currently at. targets (`Sequence[int]`): The intra-block indices `k` the loss is evaluated at, each `n <= k < N`. advance (`int`): How many intervals the carried state moves after the step, i.e. `L_min`. Returns: `torch.Tensor` of shape `(2 * len(targets) + 1, N)`: for every target, the displacement from `X_n` to `X_k` followed by the row that selects `u_k`; then, last, the displacement from `X_n` to `X_{n+L_min}`. """ plan = torch.zeros(2 * len(targets) + 1, step_sizes.shape[0], dtype=step_sizes.dtype, device=step_sizes.device) for position, target in enumerate(targets): plan[2 * position, start:target] = step_sizes[start:target] plan[2 * position + 1, target] = 1.0 plan[-1, start : start + advance] = step_sizes[start : start + advance] return plan def pdd_sampling_plan(step_sizes: torch.Tensor, start: int, block_size: int) -> torch.Tensor: r""" The single direction a PDD generation step needs: the *mean* velocity of the whole block. Normalizing the fused displacement by the block span turns it into the block's average velocity, which is what an ordinary Euler step over the block boundaries consumes — so a plain rectified-flow scheduler drives PDD generation unchanged. Args: step_sizes (`torch.Tensor` of shape `(N,)`): The grid step sizes `h_l = t_{l+1} - t_l`. start (`int`): The block start `n`. block_size (`int`): The block size `L`. Returns: `torch.Tensor` of shape `(1, N)`: the plan. """ plan = torch.zeros(1, step_sizes.shape[0], dtype=step_sizes.dtype, device=step_sizes.device) span = step_sizes[start : start + block_size].sum() plan[0, start : start + block_size] = step_sizes[start : start + block_size] / span return plan class PDDParallelHead(nn.Module): r""" The `N` per-interval output heads of a PDD parallel decoder, in place of one final linear layer. The heads are held as a single `(num_steps, out_features, in_features)` parameter, every slice initialized from the pre-trained layer this replaces — so at initialization every interval predicts exactly the teacher's velocity and the parallel decoder starts as the teacher. That pre-trained layer is also kept as a frozen buffer pair, which is what [`pdd_teacher_mode`] switches to: the teacher's instantaneous velocity stays available from the same module after the heads have moved. `forward` does not evaluate the heads one by one: it fuses them into the `num_directions` linear maps of the current `plan` and applies those, which is the paper's layer fusion (§3.1) and keeps the head's cost independent of `num_steps`. Args: source (`nn.Linear`): The pre-trained final layer to repeat. num_steps (`int`): The grid size `N`, i.e. how many heads to hold. """ def __init__(self, source: nn.Linear, num_steps: int): super().__init__() self.num_steps = num_steps self.in_features = source.in_features self.out_features = source.out_features self.weight = nn.Parameter(source.weight.detach()[None].repeat(num_steps, 1, 1).clone()) self.bias = ( None if source.bias is None else nn.Parameter(source.bias.detach()[None].repeat(num_steps, 1).clone()) ) self.register_buffer("teacher_weight", source.weight.detach().clone(), persistent=False) self.register_buffer( "teacher_bias", None if source.bias is None else source.bias.detach().clone(), persistent=False ) self.teacher = False # `plan` is a plain attribute rather than a buffer: it is per-step control flow, not model state to # serialize or shard. The default reproduces the source layer, so an unplanned head is the teacher's head. self.plan = torch.zeros(1, num_steps) self.plan[0, 0] = 1.0 def set_plan(self, plan: torch.Tensor) -> None: r""" Set the `(num_directions, num_steps)` coefficient matrix the next forward fuses the heads with. Args: plan (`torch.Tensor`): The plan. Row `p` weights the `N` heads into the `p`-th output direction. """ if plan.ndim != 2 or plan.shape[1] != self.num_steps: raise ValueError( f"A PDD plan must be a `(num_directions, {self.num_steps})` matrix, got {list(plan.shape)}." ) self.plan = plan @property def num_directions(self) -> int: return 1 if self.teacher else self.plan.shape[0] def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: r""" Args: hidden_states (`torch.Tensor` of shape `(..., in_features)`): The backbone's final hidden state. Returns: `torch.Tensor` of shape `(..., num_directions * out_features)`: the planned directions, stacked on the channel axis in plan-row order. Under [`pdd_teacher_mode`] this is the single pre-trained direction. """ if self.teacher: return F.linear(hidden_states, self.teacher_weight, self.teacher_bias) plan = self.plan.to(device=self.weight.device, dtype=self.weight.dtype) weight = torch.einsum("pn,noi->poi", plan, self.weight).flatten(0, 1) bias = None if self.bias is None else torch.einsum("pn,no->po", plan, self.bias).flatten() return F.linear(hidden_states, weight, bias) class PDDLoRALinear(nn.Module): r""" A frozen `nn.Linear` with a trainable low-rank update, `y = W x + b + (alpha / rank) * B A x`. The PDD counterpart of `lora_utils.py`'s `LoRAModule`, and deliberately not built on it: the adapter is a node in the model tree (so FSDP shards it and the base layer stays visible for dtype pinning), not an out-of-tree forward patch, and it must collapse to exactly the frozen layer when disabled. The adapter parameters are held in float32 and cast to the activation dtype inside `forward`, so the optimizer sees float32 master weights while the matmuls stay at the backbone's precision. `B` starts at zero, so the wrapped module is exactly the frozen layer at initialization — and is again exactly the frozen layer whenever `enabled` is false, which is how [`pdd_teacher_mode`] recovers the teacher without a second copy of the backbone. Args: base (`nn.Linear`): The layer to wrap. It is frozen here. rank (`int`): The rank of the update. alpha (`float`): The scaling numerator; `alpha == rank` means a unit-scaled update. """ def __init__(self, base: nn.Linear, rank: int, alpha: float): super().__init__() self.base = base self.base.requires_grad_(False) self.scaling = alpha / rank self.enabled = True self.lora_down = nn.Parameter(torch.empty(rank, base.in_features, dtype=torch.float32)) self.lora_up = nn.Parameter(torch.zeros(base.out_features, rank, dtype=torch.float32)) nn.init.kaiming_uniform_(self.lora_down, a=5**0.5) # Models may read `linear.weight.dtype` off their projections to align activations with a mixed-precision # checkpoint (MiniMax-H3 does), so the wrapper has to present the wrapped layer's own tensors under the usual # names. @property def weight(self) -> torch.Tensor: return self.base.weight @property def bias(self) -> Optional[torch.Tensor]: return self.base.bias @property def in_features(self) -> int: return self.base.in_features @property def out_features(self) -> int: return self.base.out_features def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: out = self.base(hidden_states) if not self.enabled: return out update = F.linear( F.linear(hidden_states, self.lora_down.to(hidden_states.dtype)), self.lora_up.to(hidden_states.dtype), ) return out + self.scaling * update.to(out.dtype) def add_pdd_lora(module: nn.Module, target_names: Sequence[str], rank: int, alpha: float) -> int: r""" Wrap every `nn.Linear` whose qualified name ends in one of `target_names` with a [`PDDLoRALinear`], in place. Args: module (`nn.Module`): The root to walk. target_names (`Sequence[str]`): Qualified-name suffixes to match, e.g. `("to_q", "ff.net.2")`. rank (`int`): The rank of every adapter. alpha (`float`): The scaling numerator of every adapter. Returns: `int`: The number of layers wrapped. """ targets = [ (name, child) for name, child in module.named_modules() if isinstance(child, nn.Linear) and any(name.endswith(suffix) for suffix in target_names) ] for name, child in targets: parent_name, _, attribute = name.rpartition(".") parent = module.get_submodule(parent_name) if parent_name else module setattr(parent, attribute, PDDLoRALinear(child, rank, alpha)) return len(targets) def merge_pdd_lora(module: nn.Module) -> int: r""" Fold every [`PDDLoRALinear`] into its frozen base layer and unwrap it, in place. An inference-only optimization: afterwards each wrapped layer is a plain `nn.Linear` whose weight already carries `scaling * B A`, so a forward costs one matmul instead of three. The [`PDDParallelHead`]s are left untouched — their effective weight changes every step with the `plan`, so they cannot be folded into a static layer. Two things to know before calling this. Merging overwrites the base weight, so it destroys the `enabled=False` fallback [`pdd_teacher_mode`] relies on: only merge on a pure student inference path. And the update is accumulated in float32 (the adapters' storage dtype) then cast back to the base weight's dtype, so the result is not bit-for-bit the un-merged forward — it rounds the low-rank delta into the backbone precision. Call this before any device offload / quantization / FSDP wrap is registered on the model, so no hook has to be rebuilt and the delta lands on the unquantized weight. Args: module (`nn.Module`): The root to walk, e.g. the transformer returned by `load_pdd_lora`. Returns: `int`: The number of adapters merged. """ # Collect first, then mutate: replacing a child while walking `module.modules()` would perturb the traversal. adapters = [ (parent, attribute, child) for parent in module.modules() for attribute, child in parent.named_children() if isinstance(child, PDDLoRALinear) ] for parent, attribute, adapter in adapters: base = adapter.base with torch.no_grad(): weight = base.weight delta = adapter.scaling * ( adapter.lora_up.to(weight.device) @ adapter.lora_down.to(weight.device) ) weight.data = (weight.data.to(delta.dtype) + delta).to(weight.dtype) setattr(parent, attribute, base) return len(adapters) @contextlib.contextmanager def pdd_teacher_mode(transformer): r""" Run `transformer` as the frozen pre-trained teacher. The low-rank updates of the backbone are switched off and every [`PDDParallelHead`] falls back to the weights it was built from, so the forward is bit-for-bit the released model's instantaneous velocity — with a single output direction rather than the planned ones. """ heads = [module for module in transformer.modules() if isinstance(module, PDDParallelHead)] adapters = [module for module in transformer.modules() if isinstance(module, PDDLoRALinear)] for head in heads: head.teacher = True for adapter in adapters: adapter.enabled = False try: yield transformer finally: for head in heads: head.teacher = False for adapter in adapters: adapter.enabled = True def _strip_fsdp_wrapper(name: str) -> str: r"""Drop the `_fsdp_wrapped_module` path segments FSDP injects around each separately-wrapped child unit.""" return ".".join(part for part in name.split(".") if part != "_fsdp_wrapped_module") def pdd_state_dict(transformer, state_dict: Optional[dict] = None) -> dict: r""" The trainable PDD state of a parallel decoder: the enlarged heads and every low-rank update. The frozen backbone is not included, so a checkpoint is a few gigabytes rather than the full size of the base model. Args: transformer (`nn.Module`): The parallel decoder; its module tree decides which keys are trainable. state_dict (`dict`, optional): The weights to filter, defaulting to `transformer.state_dict()`. Under FSDP, pass an already-gathered `FULL_STATE_DICT` here, since the live module views are sharded. FSDP wraps every child unit in a `_fsdp_wrapped_module`, so a wrapped module's `named_modules` path (`blocks.0._fsdp_wrapped_module.attn.to_q`) never matches the clean keys of a gathered `FULL_STATE_DICT` (`blocks.0.attn.to_q.lora_down`) and only the root wrap unit would survive the filter. Both sides are normalized through [`_strip_fsdp_wrapper`] so the block LoRA and the parallel heads are kept too. """ if state_dict is None: state_dict = transformer.state_dict() trainable = { _strip_fsdp_wrapper(name) for name, module in transformer.named_modules() if isinstance(module, (PDDParallelHead, PDDLoRALinear)) } return { _strip_fsdp_wrapper(name): value.detach().cpu() for name, value in state_dict.items() if any(_strip_fsdp_wrapper(name).startswith(f"{prefix}.") for prefix in trainable) and ".base." not in name } PDD_WEIGHTS_NAME = "pdd.safetensors" PDD_EMA_WEIGHTS_NAME = "pdd_ema.safetensors" # Pre-rename checkpoints stored live weights here; resume still accepts it. PDD_LEGACY_LIVE_WEIGHTS_NAME = "pdd_live.safetensors" # The released MiniMax-H3 recipe. Every field is meant to be overridden by the `pdd_config.json` a training run # writes next to its weights; a port to another model passes its own `defaults` to [`load_pdd_config`] instead. PDD_DEFAULT_CONFIG = { "pdd_num_steps": 32, "pdd_block_size": 4, "lora_rank": 64, "lora_alpha": 64.0, "lora_targets": "to_q,to_k,to_v,to_out.0,ff.net.0.proj,ff.net.2,adaln_proj.linear", } def resolve_pdd_lora_path(path): r""" A checkpoint directory or a weights file. A directory prefers `pdd_ema.safetensors` (the EMA inference export) and falls back to `pdd.safetensors` (live weights, or the EMA file on checkpoints written before the rename). """ if path is None: return None path = os.path.abspath(os.path.expanduser(path)) if os.path.isdir(path): ema = os.path.join(path, PDD_EMA_WEIGHTS_NAME) live = os.path.join(path, PDD_WEIGHTS_NAME) if os.path.isfile(ema): path = ema elif os.path.isfile(live): path = live else: raise FileNotFoundError( f"PDD checkpoint directory {path} has neither {PDD_EMA_WEIGHTS_NAME} nor {PDD_WEIGHTS_NAME}." ) if not os.path.isfile(path): raise FileNotFoundError(f"PDD checkpoint does not exist: {path}") return path def load_pdd_config(weights_path, defaults=None): r"""Rank / alpha / targets / grid next to the weights file (`pdd_config.json`), over `defaults`.""" config = dict(PDD_DEFAULT_CONFIG if defaults is None else defaults) config_path = os.path.join(os.path.dirname(weights_path), "pdd_config.json") if os.path.isfile(config_path): with open(config_path, encoding="utf-8") as handle: saved = json.load(handle) aliases = {"lora_rank": "rank", "lora_alpha": "network_alpha", "lora_targets": "target_name"} for key in config: if key in saved: config[key] = saved[key] elif aliases.get(key) in saved: config[key] = saved[aliases[key]] if not isinstance(config["lora_targets"], str): config["lora_targets"] = ",".join(config["lora_targets"]) return config def pdd_num_inference_steps(config, num_inference_steps, teacher_default=None): r"""Keep `num_inference_steps` when it divides `N`; otherwise snap a leftover teacher default to `N / L`.""" grid = int(config["pdd_num_steps"]) steps = int(num_inference_steps) if grid % steps == 0: return steps block = int(config["pdd_block_size"]) if teacher_default is not None and steps == int(teacher_default) and block > 0 and grid % block == 0: nfe = grid // block print(f"PDD checkpoint: using num_inference_steps {nfe} (grid {grid}, block {block})", flush=True) return nfe raise ValueError(f"num_inference_steps {steps} must divide PDD grid size {grid}.")