diff --git a/README.md b/README.md index df0e6201..917f5d64 100644 --- a/README.md +++ b/README.md @@ -46,7 +46,7 @@ Each model links to its recipe page; every documented recipe is labeled with a [ | **[Wan2.2-T2V-A14B](https://miles.radixark.com/docs/diffusion/models/wan/wan2-2)** | T2V | Flow-GRPO + PickScore, LoRA SFT | | **[LTX-2.3](https://miles.radixark.com/docs/diffusion/models/ltx/ltx2)** | T2V | Flow-GRPO + PickScore | | **[Cosmos3 (Edge / Nano / Super)](https://miles.radixark.com/docs/diffusion/models/cosmos/cosmos3)** | T2I | Flow-GRPO + PickScore | -| **[MiniMax H3](https://miles.radixark.com/docs/diffusion/models/h3/h3)** | T2VA | Open ([PR #154](https://github.com/radixark/miles_diffusion/pull/154)); 2-GPU recipe, verified; large-scale coming soon | +| **[MiniMax H3](https://miles.radixark.com/docs/diffusion/models/h3/h3)** | T2VA | Flow-GRPO + PickScore (t2va, video branch); 2-GPU recipe, verified | --- diff --git a/docs/index.md b/docs/index.md index 64dfd53e..591b820a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -45,7 +45,7 @@ appear in the [Miles model list](https://miles.radixark.com/docs#supported-model | [Wan2.2-T2V-A14B](models/wan/wan2-2.md) | T2V | Flow-GRPO + PickScore, LoRA SFT | | [LTX-2.3](models/ltx/ltx2.md) | T2V | Flow-GRPO + PickScore | | [Cosmos3-Nano](models/cosmos/cosmos3.md) | T2I | Flow-GRPO + PickScore | -| [MiniMax H3](models/h3/h3.md) | T2VA | **Not merged** — [PR #154](https://github.com/radixark/miles_diffusion/pull/154); 2-GPU PR-only recipe | +| [MiniMax H3](models/h3/h3.md) | T2VA | Flow-GRPO + PickScore (t2va, video branch); 2-GPU recipe | diff --git a/docs/models/h3/h3.md b/docs/models/h3/h3.md index 34243bef..60b32b96 100644 --- a/docs/models/h3/h3.md +++ b/docs/models/h3/h3.md @@ -1,102 +1,83 @@ --- title: MiniMax H3 -description: Text-to-audio-video Flow-GRPO + PickScore on MiniMax H3 — 2-GPU recipe (not merged yet). +description: Text-to-audio-video Flow-GRPO + PickScore on MiniMax H3 — the verified 2-GPU recipe. --- -> **Not merged yet.** H3 support is only in -> [miles-diffusion#154](https://github.com/radixark/miles_diffusion/pull/154) and -> [sglang-diffusion#34365](https://github.com/sgl-project/sglang/pull/34365). -> Neither repo has it on `main`. After a normal install, follow [Setup](#3-setup) to -> move both checkouts onto those PRs. - ## 1. Model introduction [MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) is MiniMax's omni-modal generative system: one transformer denoises a packed text + video + audio sequence and emits video with native stereo audio. MiniMax distilled CFG into the checkpoint, so the official setting has **no classifier-free guidance** — the recipe is unguided -(`--diffusion-guidance-scale 1.0`). The unmerged PR trains the **t2va** (text-to-audio-video) workflow, video +(`--diffusion-guidance-scale 1.0`). The recipe trains the **t2va** (text-to-audio-video) workflow, video branch only — the audio stream is rolled out but not trained. ## 2. Supported variants | Model | HF ID | Notes | |---|---|---| -| MiniMax-H3 | [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3) | t2va video-only Flow-GRPO on [PR #154](https://github.com/radixark/miles_diffusion/pull/154). `fl2va` / `ref2va` are not trained yet. | +| MiniMax-H3 | [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3) | t2va video-only Flow-GRPO. `fl2va` / `ref2va` are not trained yet. | -## 3. Setup +## 3. Family config -H3 is not on `main` of either repo. Finish [Installation](../../getting-started/installation.md) first -(Docker image or from source), then move **both** halves onto the open PRs — a `main` trainer -against a `main` engine will not run this recipe. +From `miles/backends/fsdp_utils/configs/h3.py`: -| Repo | PR | What it adds | +| Property | Value | Why | |---|---|---| -| miles-diffusion | [radixark/miles_diffusion#154](https://github.com/radixark/miles_diffusion/pull/154) | H3 family config, LoRA IPC map, 2-GPU recipe | -| sglang-diffusion | [sgl-project/sglang#34365](https://github.com/sgl-project/sglang/pull/34365) | MiniMax H3 `t2va` RL rollout | - -Paths below match the Docker image (`/root/miles_diffusion`, `/sgl-workspace/sglang`). - -```bash -# miles-diffusion → PR 154 -cd /root/miles_diffusion -git fetch https://github.com/radixark/miles_diffusion.git pull/154/head:pr-154 -git checkout pr-154 -pip install --force-reinstall -r requirements.txt # H3 needs the diffusers git pin -pip install -e . --no-deps - -# sglang-diffusion → PR 34365 -cd /sgl-workspace/sglang -git fetch https://github.com/sgl-project/sglang.git pull/34365/head:pr-34365 -git checkout pr-34365 -pip install -e "python[all]" -``` - -Drop these checkouts once both PRs merge. +| CFG training | Off (asserted) | CFG is distilled into the checkpoint; the forward is unguided | +| `sde_timestep_divisor` | 1000.0 | The DiT conditions on `1 - sigma`, so the recorded timestep is divided back down | +| Weight sync | `--use-lora --lora-ipc-weight-sync` (asserted) | sgl-d's H3 DiT renames modules and fuses Q/K/V; other sync modes push names the engine drops with a warning, silently training nothing | +| LoRA targets | `attn.to_{q,k,v}`, `attn.to_out.0`, `ff.net.0.proj`, `ff.net.2` | Grouped for the fused engine layers by `h3_weight_key_mapper.collect_h3_lora_layer_groups` | +| Optimizer state | `audio` allowed missing | The audio branch is rolled out but never trained | +| Sample micro-batch | 1 (asserted) | One packed sequence per forward | +| Forced sampling params | `task=t2va`, `short_edge=768`, `conditions=[]` | sgl-d accepts only these for H3, so none of them is exposed as an argument | ## 4. Launch -The 2-GPU recipe lives on [PR #154](https://github.com/radixark/miles_diffusion/pull/154), not `main`: -`scripts/run_diffusion_grpo_h3_t2va_2gpu.py` — 2 colocate GPUs, reward colocated, 16:9 / 4 s -(sglang pins `short_edge=768`, so the canvas is 1344×768 / 107 frames). Requires `--use-lora --lora-ipc-weight-sync` -(the recipe already sets both). +Canonical recipe: `scripts/run_diffusion_grpo_h3_t2va_2gpu.py` — train, rollout, and PickScore +colocated on 2 GPUs; 16:9 / 4 s (with `short_edge=768` the canvas is 1344×768 / 107 frames). -**Status:** [📈 V — Verified](../../user-guide/recipe-verification.md#v), based on the complete run reported on PR #154. -The command does not exist on `main`. +**Status:** [📈 V — Verified](../../user-guide/recipe-verification.md#v) ```bash -# alignment check (optimizer frozen) +# alignment diagnostic: freeze the weights, so log_prob_mean_abs_diff measures +# train-vs-rollout deviation rather than parameter drift python3 scripts/run_diffusion_grpo_h3_t2va_2gpu.py \ - --cuda-visible-devices 0,1 \ - --num-rollout 2 --eval-interval 0 --debug-alignment + --num-rollout 2 --eval-interval 0 --extra-args "--debug-skip-optimizer-step" -# short train +# long run at the resource-limited scale python3 scripts/run_diffusion_grpo_h3_t2va_2gpu.py \ - --cuda-visible-devices 0,1 --num-rollout 5 --eval-interval 0 - -# 2-GPU long run (resource-limited scale) -python3 scripts/run_diffusion_grpo_h3_t2va_2gpu.py \ - --rollout-batch-size 4 --n-samples-per-prompt 16 \ - --cuda-visible-devices 0,1 + --rollout-batch-size 4 --n-samples-per-prompt 16 ``` -The commands above are the 2-GPU verification path on -[PR #154](https://github.com/radixark/miles_diffusion/pull/154). No larger canonical H3 topology is committed to -`main`. +## 5. Recipe notes + +H3 solves `x0 = x + sigma * v`, the opposite sign from the flow-matching convention the shared +train and rollout steps assume, so `compute_noise_pred` returns the negated velocity. Both halves +must agree: a flipped sign still produces plausible latents and trains on garbage. -## 5. Reference results +The rollout engine needs sglang's `sglang-miles-h3` branch, which +[Installation](../../getting-started/installation.md) already builds from +(`SGLANG_DIFFUSION_BRANCH`). An engine built from sglang `main` rejects the run with +`MiniMax H3 does not support rollout`. -2-GPU long run from [PR #154](https://github.com/radixark/miles_diffusion/pull/154) -(`rollout_batch_size=4`, `n_samples_per_prompt=16`): +## 6. Reference results + +2-GPU long run (`rollout_batch_size=4`, `n_samples_per_prompt=16`): - `train/log_prob_mean_abs_diff` stays in **1.4e-5 .. 7.1e-5** (under `1e-4`). - Eval PickScore (`eval/pickscore_test`, every 10 rollouts): **0.806 → 0.814**, peak **0.816** at step 59. - Train `rollout/reward/raw_mean` sits around **~0.79** (64 samples / step; noisier than eval). -This is the batch-4 curve reported on [PR #154](https://github.com/radixark/miles_diffusion/pull/154); it is not evidence -for a larger topology. +This is the batch-4 curve; it is not evidence for a larger topology, and no larger canonical H3 +topology is committed yet. ![MiniMax H3 log_prob_mean_abs_diff](../../assets/images/h3/log_prob_diff.png) ![MiniMax H3 eval PickScore](../../assets/images/h3/reward_eval.png) ![MiniMax H3 train rollout PickScore](../../assets/images/h3/reward_rollout.png) + +## 7. Pairs well with + +- [LoRA Training and Weight Sync](../../advanced/lora.md) — the IPC merge this recipe requires. +- [Rewards](../../user-guide/rewards.md) — PickScore worker pool configuration. diff --git a/docs/user-guide/recipe-verification.md b/docs/user-guide/recipe-verification.md index 24f0dc4c..4a4ee768 100644 --- a/docs/user-guide/recipe-verification.md +++ b/docs/user-guide/recipe-verification.md @@ -46,8 +46,7 @@ count as verification. Flow-GRPO + PickScore. - `run_diffusion_grpo_cosmos3_pickscore_t2i_4gpu.py` — Cosmos3-Nano Flow-GRPO + PickScore. - - `run_diffusion_grpo_h3_t2va_2gpu.py` — MiniMax H3 Flow-GRPO + PickScore - ([PR #154](https://github.com/radixark/miles_diffusion/pull/154) only; not on `main`). + - `run_diffusion_grpo_h3_t2va_2gpu.py` — MiniMax H3 Flow-GRPO + PickScore. - **○ NV** - `run_diffusion_grpo_wan22_pickscore_5gpu.py` — Wan2.2 5-GPU LoRA Flow-GRPO + PickScore. diff --git a/miles/backends/fsdp_utils/configs/h3.py b/miles/backends/fsdp_utils/configs/h3.py new file mode 100644 index 00000000..5912997c --- /dev/null +++ b/miles/backends/fsdp_utils/configs/h3.py @@ -0,0 +1,185 @@ +"""MiniMax H3 family config: t2va video-only Flow-GRPO.""" + +from __future__ import annotations + +from argparse import Namespace + +import torch + +from miles.utils.types import CondKwargs + +from .train_pipeline_config import TrainPipelineConfig, register_train_pipeline_config + +AUDIO_IN_CHANNELS = 32 + + +@register_train_pipeline_config("h3") +class H3TrainPipelineConfig(TrainPipelineConfig): + """MiniMax H3 t2va video-only GRPO (audio branch frozen / deterministic in rollout).""" + + hf_ckpt_name_patterns = ("minimax-h3", "minimax_h3", "/h3") + supports_cfg_training = False + sde_timestep_divisor = 1000.0 + optimizer_state_allowed_missing = ["audio"] + lora_layer_group_collector_path = "miles.backends.fsdp_utils.h3_weight_key_mapper.collect_h3_lora_layer_groups" + + lora_target_modules = [ + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "ff.net.0.proj", + "ff.net.2", + ] + + @classmethod + def validate_args(cls, args: Namespace) -> None: + # sglang's H3 DiT renames modules and fuses Q/K/V, so weights only reach the + # rollout through the LoRA IPC path's layer grouper; any other sync mode would + # push names the engine drops with a warning, silently training nothing. + if not (args.use_lora and args.lora_ipc_weight_sync): + raise ValueError("H3 training requires --use-lora with --lora-ipc-weight-sync") + + @classmethod + def apply_rollout_sampling_params( + cls, + args: Namespace, + sampling_params: dict, + extra_sampling_params: dict, + ) -> None: + extra_sampling_params.update( + { + # sgl-d accepts only task=t2va for rollout and short_edge=768 for any + # H3 request, so neither is exposed as an argument. + "task": "t2va", + "conditions": [], + "target": { + "short_edge": 768, + "aspect_ratio": str(args.diffusion_h3_aspect_ratio), + "duration_seconds": float(args.diffusion_h3_duration_seconds), + }, + "audio_flow_shift": float(args.diffusion_audio_flow_shift), + } + ) + if args.diffusion_flow_shift is not None: + extra_sampling_params["flow_shift"] = float(args.diffusion_flow_shift) + # MiniMaxH3SamplingParams marks CFG/canvas fields init=False; canvas comes from target. + extra_sampling_params.pop("guidance_scale_2", None) + for key in ( + "guidance_scale", + "guidance_scale_2", + "true_cfg_scale", + "negative_prompt", + "width", + "height", + "num_frames", + "fps", + ): + sampling_params.pop(key, None) + + def prepare_cond_kwargs(self, cond: CondKwargs | None, device: torch.device) -> dict: + if cond is None: + return {} + kwargs: dict = {} + if cond.encoder_hidden_states: + enc = torch.cat(cond.encoder_hidden_states).to(device) + if enc.ndim == 2: + enc = enc.unsqueeze(0) + kwargs["encoder_hidden_states"] = enc + if cond.h3_packed_layout is not None: + kwargs["h3_packed_layout"] = { + k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in cond.h3_packed_layout.items() + } + if cond.h3_token_tags is not None: + kwargs["h3_token_tags"] = cond.h3_token_tags.to(device) + return kwargs + + def collate_cond_for_sample_batch( + self, + per_sample_cond_kwargs: list[dict], + device: torch.device, + pad_to_len: int | None = None, + ) -> dict: + if len(per_sample_cond_kwargs) != 1: + raise NotImplementedError("H3 GRPO currently requires micro-batch-size-sample=1") + return {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in per_sample_cond_kwargs[0].items()} + + def compute_noise_pred( + self, + *, + model: torch.nn.Module, + latents_input: torch.Tensor, + timesteps_input: torch.Tensor, + pos_cond: dict | None, + neg_cond: dict | None, + joint_cond: dict | None, + use_cfg: bool, + cfg_batching: bool, + guidance_scale: float, + true_cfg_scale: float | None, + ) -> torch.Tensor: + del neg_cond, joint_cond, use_cfg, cfg_batching, guidance_scale, true_cfg_scale + cond = dict(pos_cond or {}) + packed = cond.get("h3_packed_layout") + token_tags = cond.get("h3_token_tags") + encoder_hidden_states = cond.get("encoder_hidden_states") + if packed is None or token_tags is None or encoder_hidden_states is None: + raise ValueError("H3 train requires h3_packed_layout, h3_token_tags, encoder_hidden_states in pos_cond") + + device = latents_input.device + dtype = latents_input.dtype + + # latents_input: [B, num_video_target_rows, width] + bsz = latents_input.shape[0] + if bsz != 1: + raise NotImplementedError("H3 packed forward supports batch size 1 for now") + + layout = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in packed.items()} + tags = ( + token_tags.to(device) if isinstance(token_tags, torch.Tensor) else torch.tensor(token_tags, device=device) + ) + sigma = (timesteps_input.float() / float(self.sde_timestep_divisor)).view(-1) + timestep = 1.0 - sigma + seq_len = int(layout["seq_len"]) + width = latents_input.shape[-1] + + img_pos = layout["img_pos"].view(-1).long().to(device) + audio_pos = layout["audio_pos"].view(-1).long().to(device) + update_mask = layout["update_mask"].view(-1).bool().to(device) + text_pos = layout["text_pos"].view(-1).long().to(device) + + # The transformer takes one row block per modality, each ordered like its + # ``*_indices``, and scatters them into the packed buffer itself. Only the + # target rows are replayed; conditioning rows stay zero, as does the audio + # stream (H3 GRPO trains the video branch only). + video_hidden = torch.zeros(1, int(img_pos.shape[0]), width, device=device, dtype=dtype) + video_hidden[0, update_mask] = latents_input[0].to(dtype) + audio_hidden = torch.zeros(1, int(audio_pos.shape[0]), AUDIO_IN_CHANNELS, device=device, dtype=dtype) + + out = model( + hidden_states=video_hidden, + audio_hidden_states=audio_hidden, + encoder_hidden_states=encoder_hidden_states.to(dtype), + timestep=timestep.to(dtype), + timestep_indices=layout.get("timestep_indices", torch.zeros(seq_len, device=device, dtype=torch.long)), + # Padding rows carry tag -1; the AdaLN table is indexed by tag, so they + # must be folded onto a real modality exactly as the rollout does. + token_tags=tags.long().clamp(min=0), + position_ids=layout["img_position_ids"].to(device=device, dtype=torch.float32), + video_indices=img_pos, + audio_indices=audio_pos, + text_indices=text_pos, + ) + velocity = out[0] if isinstance(out, tuple) else out.sample + # Rows follow video_indices; keep the target subset and return the + # diffusers-compatible flow direction (negated H3 velocity). + return (-velocity[0, update_mask]).to(dtype) + + def cfg_combine( + self, + noise_pred_pos: torch.Tensor, + noise_pred_neg: torch.Tensor, + guidance_scale: float, + true_cfg_scale: float | None = None, + ) -> torch.Tensor: + raise NotImplementedError("H3 distilled CFG into the checkpoint; the forward is unguided") diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index e55df3fe..dec07c3f 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -66,6 +66,7 @@ def resolve_diffusion_model_family(model_ref: str) -> str: def get_train_pipeline_config_cls(family: str) -> type[TrainPipelineConfig]: """The TrainPipelineConfig class registered for a resolved family key.""" + _populate_registry() cls = _REGISTRY.get(family.lower()) if cls is None: raise ValueError( @@ -96,7 +97,25 @@ class TrainPipelineConfig(abc.ABC): def validate_args(cls, args) -> None: """Family-specific arg validation/defaults; runs once at arg validation.""" + @classmethod # noqa: B027 — optional hook, deliberately non-abstract + def apply_rollout_sampling_params( + cls, + args, + sampling_params: dict, + extra_sampling_params: dict, + ) -> None: + """Adjust the ``POST /rollout/generate`` body for this family; default adjusts nothing. + + Both dicts are mutated in place: ``sampling_params`` is the request body the + generic builder has already filled in, ``extra_sampling_params`` the + family-specific passthrough. Families whose sgl-d request schema rejects some + of the generic fields drop them here. + """ + sde_timestep_divisor = 1.0 + # LoRA IPC layer grouper for families whose rollout module names or tensor layout + # differ from the trained diffusers ones; None keeps the generic PEFT grouping. + lora_layer_group_collector_path: str | None = None def configure(self, args) -> None: # noqa: B027 optional no-op hook, not abstract """Bind the request constants a family needs at train time; default binds none.""" diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index 99beb915..ff45c213 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -456,10 +456,39 @@ def _prepare_lora_param(self, param: torch.Tensor) -> torch.Tensor: ).to_local() return param + def _collect_layer_groups( + self, model: torch.nn.Module + ) -> tuple[list[list[tuple[str, torch.Tensor]]], list[str], int]: + """Group this model's LoRA tensors into rollout layer names, per model family.""" + from miles.utils.misc import load_function + + collector_path = None + if self.args.train_pipeline_config_path: + collector_path = load_function(self.args.train_pipeline_config_path).lora_layer_group_collector_path + if collector_path is None: + return collect_lora_layer_groups(model.state_dict()) + + # A family whose rollout fuses several projections into one layer (H3's + # qkv_proj) combines adapters here, so DTensor shards must resolve first. + lora_state = { + name: self._prepare_lora_param(param) + for name, param in model.state_dict().items() + if PeftLoRAKeyMapper.is_lora_key(name) + } + layer_groups, unmapped_keys, num_lora_keys = load_function(collector_path)(lora_state) + if unmapped_keys: + # The rollout only warns about a name it cannot resolve, which would + # leave that adapter frozen at its checkpoint value. + raise ValueError( + f"{collector_path} could not map {len(unmapped_keys)} adapter modules to " + f"rollout layer names (first 5: {unmapped_keys[:5]})" + ) + return layer_groups, unmapped_keys, num_lora_keys + def update_weights(self) -> None: self.weight_version += 1 for target_module, model in self.models.items(): - layer_groups, unmapped_keys, num_lora_keys = collect_lora_layer_groups(model.state_dict()) + layer_groups, unmapped_keys, num_lora_keys = self._collect_layer_groups(model) bucket: list[tuple[str, torch.Tensor]] = [] bucket_size = 0 num_buckets = 0 @@ -490,7 +519,10 @@ def update_weights(self) -> None: num_buckets += 1 if self.weight_version <= 2 and dist.is_initialized() and dist.get_rank() == 0: - _, num_layers, sample_layers, _ = PeftLoRAKeyMapper.summarize_mapping(model.state_dict()) + # Report the layers actually pushed: a family that fuses projections + # into one rollout layer has fewer layers than PEFT modules. + num_layers = len(layer_groups) + sample_layers = [PeftLoRAKeyMapper.layer_prefix(group[0][0]) for group in layer_groups[:3]] logger.info( "LoRA IPC weight sync v%s [%s]: pushed %d lora tensors, " "%d layer prefixes in %d buckets (unmapped=%d)", diff --git a/miles/backends/fsdp_utils/h3_weight_key_mapper.py b/miles/backends/fsdp_utils/h3_weight_key_mapper.py new file mode 100644 index 00000000..764e0af7 --- /dev/null +++ b/miles/backends/fsdp_utils/h3_weight_key_mapper.py @@ -0,0 +1,182 @@ +"""Map diffusers MiniMax H3 LoRA names to sglang H3 DiT layer names. + +Training uses diffusers ``MiniMaxH3Transformer3DModel`` (separate Q/K/V). +Rollout uses sglang ``MiniMaxH3DiTModel`` (fused ``qkv_proj``). LoRA IPC sync +must therefore rename modules and stack the Q/K/V adapters before the push. + +Names that resolve to no sglang layer are skipped with a warning on the rollout +side, so an incomplete map silently freezes those adapters at their checkpoint +values — anything unrecognized is reported as unmapped instead of guessed. +""" + +from __future__ import annotations + +import re +from collections.abc import Mapping + +import torch + +_QKV_RE = re.compile( + r"^(?P(?:token_refiner\.)?refiner_blocks\.(?P\d+)|transformer_blocks\.(?P\d+))" + r"\.attn\.to_(?Pq|k|v)\.weight$" +) +# After PEFT strip, refiner path is token_refiner.refiner_blocks -> token_refiner.blocks +_QKV_RE_SGL = re.compile( + r"^(?P(?:token_refiner\.)?blocks\.(?P\d+)|blocks\.(?P\d+))" + r"\.attn\.to_(?Pq|k|v)\.weight$" +) + + +def _qkv_group_key(name: str) -> tuple[str, str] | None: + for regex in (_QKV_RE, _QKV_RE_SGL): + m = regex.match(name) + if m is None: + continue + prefix = m.group("prefix") + if prefix.startswith("token_refiner."): + block_idx = m.group("idx") + sgld_prefix = f"token_refiner.blocks.{block_idx}" + elif prefix.startswith("refiner_blocks."): + block_idx = m.group("idx") + sgld_prefix = f"token_refiner.blocks.{block_idx}" + else: + block_idx = m.group("idx2") + sgld_prefix = f"blocks.{block_idx}" + return sgld_prefix, m.group("which") + return None + + +def _swap_gated_ffn_halves(tensor: torch.Tensor) -> torch.Tensor: + """Reorder a gated FFN input projection from diffusers' halves to sglang's. + + diffusers' GEGLU splits the fused projection as ``[up, gate]`` and computes + ``up * gelu(gate)``; sglang's ``mlp.fc1`` splits it as ``[gate, up]`` and + computes ``silu(gate) * up``. Same weights, opposite halves. + """ + rows = tensor.shape[0] + if rows % 2: + raise ValueError(f"H3 gated FFN projection must have an even row count, got {rows}") + half = rows // 2 + return torch.cat([tensor[half:], tensor[:half]], dim=0) + + +_LORA_AB_RE = re.compile(r"\.lora_([AB])(?:\.[^.]+)?(?:\.weight)?$") +_PEFT_PREFIX = "base_model.model." + +# LoRA-able H3 submodules other than Q/K/V, as (diffusers suffix, sglang suffix). +# Kept as an explicit whitelist: an unrecognized module must surface as unmapped +# rather than reach the rollout under a guessed name, where it would be skipped +# with only a warning and silently freeze that adapter. +_LORA_MODULE_SUFFIXES: tuple[tuple[str, str], ...] = ( + (".attn.to_out.0", ".attn.out_proj"), + (".ff.net.0.proj", ".mlp.fc1"), + (".ff.net.2", ".mlp.fc2"), +) + +_BLOCK_PREFIX_REPLACEMENTS: tuple[tuple[str, str], ...] = ( + (r"^token_refiner\.refiner_blocks\.", "token_refiner.blocks."), + (r"^refiner_blocks\.", "token_refiner.blocks."), + (r"^transformer_blocks\.", "blocks."), +) + + +def _strip_peft_prefix(name: str) -> str: + return name[len(_PEFT_PREFIX) :] if name.startswith(_PEFT_PREFIX) else name + + +def _normalize_block_prefix(module_path: str) -> str: + out = module_path + for pattern, repl in _BLOCK_PREFIX_REPLACEMENTS: + out = re.sub(pattern, repl, out) + return out + + +def _stack_qkv_lora(triple: dict[str, dict[str, torch.Tensor]], layer: str) -> tuple[torch.Tensor, torch.Tensor]: + """Stack per-projection LoRA into the 3D layout sglang's fused qkv expects. + + ``MergedColumnParallelLinearWithLoRA`` multiplies a 3D ``B @ A`` batchwise and + flattens the result, so stacking along a leading axis yields a delta ordered + ``[q_all, k_all, v_all]`` — exactly how sglang stores ``qkv_proj.weight``. + Note this differs from the dense path, which must instead emit the head-major + grouped layout because it goes through the checkpoint weight loader. + """ + missing = {"q", "k", "v"} - set(triple) + if missing: + raise ValueError(f"H3 LoRA IPC incomplete QKV for {layer}: missing {sorted(missing)}") + order = ("q", "k", "v") + a_shapes = {triple[w]["A"].shape for w in order} + b_shapes = {triple[w]["B"].shape for w in order} + if len(a_shapes) != 1 or len(b_shapes) != 1: + raise ValueError( + f"H3 LoRA IPC expects MHA-shaped Q/K/V adapters for {layer}, " + f"got A={sorted(a_shapes)} B={sorted(b_shapes)}" + ) + lora_a = torch.stack([triple[w]["A"] for w in order], dim=0) + lora_b = torch.stack([triple[w]["B"] for w in order], dim=0) + return lora_a, lora_b + + +def collect_h3_lora_layer_groups( + state_dict: Mapping[str, torch.Tensor], +) -> tuple[list[list[tuple[str, torch.Tensor]]], list[str], int]: + """Group PEFT LoRA tensors into sglang H3 layer names for IPC weight sync. + + Returns ``(layer_groups, unmapped_keys, num_lora_keys)`` to match + ``collect_lora_layer_groups``; each group holds one layer's A/B pair so they + always land in the same IPC bucket. + """ + per_module: dict[str, dict[str, torch.Tensor]] = {} + unmapped: list[str] = [] + num_lora_keys = 0 + + for name, tensor in state_dict.items(): + if ".lora_A" not in name and ".lora_B" not in name: + continue + stripped = _strip_peft_prefix(name) + match = _LORA_AB_RE.search(stripped) + if match is None: + unmapped.append(name) + continue + per_module.setdefault(stripped[: match.start()], {})[match.group(1)] = tensor + num_lora_keys += 1 + + qkv_pending: dict[str, dict[str, dict[str, torch.Tensor]]] = {} + simple: dict[str, dict[str, torch.Tensor]] = {} + + for module_path, ab in per_module.items(): + if "A" not in ab or "B" not in ab: + unmapped.append(module_path) + continue + + probe = f"{module_path}.weight" + qkv = _qkv_group_key(probe) + if qkv is not None: + sgld_prefix, which = qkv + qkv_pending.setdefault(f"{sgld_prefix}.attn.qkv_proj", {})[which] = ab + continue + + normalized = _normalize_block_prefix(module_path) + for diffusers_suffix, sglang_suffix in _LORA_MODULE_SUFFIXES: + if normalized.endswith(diffusers_suffix): + layer = normalized[: -len(diffusers_suffix)] + sglang_suffix + if diffusers_suffix == ".ff.net.0.proj": + # Gated FFN: diffusers stores [up, gate], sglang fc1 wants + # [gate, up]. Only B is row-indexed by output, so A is untouched. + ab = {"A": ab["A"], "B": _swap_gated_ffn_halves(ab["B"])} + simple[layer] = ab + break + else: + unmapped.append(module_path) + + groups: list[list[tuple[str, torch.Tensor]]] = [] + for layer in sorted(simple): + ab = simple[layer] + groups.append([(f"{layer}.lora_A", ab["A"]), (f"{layer}.lora_B", ab["B"])]) + for layer in sorted(qkv_pending): + lora_a, lora_b = _stack_qkv_lora(qkv_pending[layer], layer) + groups.append([(f"{layer}.lora_A", lora_a), (f"{layer}.lora_B", lora_b)]) + + return groups, unmapped, num_lora_keys + + +__all__ = ["collect_h3_lora_layer_groups"] diff --git a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py index 8e54fbd5..ee49bbd1 100644 --- a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py +++ b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py @@ -24,7 +24,6 @@ def prepare_flow_grpo_batch( args = ctx.args device = ctx.device config = ctx.train_pipeline_config - num_train_timesteps = int(ctx.scheduler.config.num_train_timesteps) bsz = len(batch) latents = _stack_pair_field(batch, "latent", device) @@ -47,6 +46,9 @@ def prepare_flow_grpo_batch( if len(ctx.models) == 1: component_name, model = next(iter(ctx.models.items())) else: + # Only multi-DiT families route timesteps to a phase, and only their schedulers are + # guaranteed to declare num_train_timesteps. + num_train_timesteps = int(ctx.scheduler.config.num_train_timesteps) components = {config.component_for_timestep(t, num_train_timesteps) for t in timesteps.tolist()} if len(components) > 1: raise ValueError( diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index 353219bc..b1248a24 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/miles/backends/fsdp_utils/model_backend.py @@ -246,9 +246,11 @@ def _resolve_component_class(cls, args, component: str): @staticmethod def _component_class(spec): - if not isinstance(spec, (list, tuple)) or len(spec) != 2: + # Classic pipelines write ``[library, class_name]``; modular ones append a + # loading-spec dict (H3's MiniMaxH3ModularPipeline). Nothing else is valid. + if not isinstance(spec, (list, tuple)) or len(spec) not in (2, 3): return None - library, class_name = spec + library, class_name = spec[0], spec[1] if not library or not class_name: return None try: diff --git a/miles/backends/fsdp_utils/models/diffusers/h3/__init__.py b/miles/backends/fsdp_utils/models/diffusers/h3/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/miles/backends/fsdp_utils/models/diffusers/h3/__init__.py @@ -0,0 +1 @@ + diff --git a/miles/backends/fsdp_utils/models/diffusers/h3/parallel_plan.py b/miles/backends/fsdp_utils/models/diffusers/h3/parallel_plan.py new file mode 100644 index 00000000..6f539c9c --- /dev/null +++ b/miles/backends/fsdp_utils/models/diffusers/h3/parallel_plan.py @@ -0,0 +1,6 @@ +from miles.backends.fsdp_utils.models.parallel_plan import FSDPParallelPlan + + +# Empty: MiniMaxH3Transformer3DModel already declares its _no_split_modules, and no +# submodule needs to be pinned off the global param dtype. +FSDP_PARALLEL_PLAN = FSDPParallelPlan() diff --git a/miles/rollout/sglang_diffusion_rollout.py b/miles/rollout/sglang_diffusion_rollout.py index 72c91377..cde28da0 100644 --- a/miles/rollout/sglang_diffusion_rollout.py +++ b/miles/rollout/sglang_diffusion_rollout.py @@ -71,10 +71,15 @@ def build_rollout_sampling_params( if args.diffusion_fps is not None: sampling_params["fps"] = int(args.diffusion_fps) + extra_sampling_params = dict(extra_sampling_params or {}) if args.diffusion_guidance_scale_2 is not None: - extra_sampling_params = dict(extra_sampling_params or {}) extra_sampling_params["guidance_scale_2"] = float(args.diffusion_guidance_scale_2) + if args.train_pipeline_config_path: + load_function(args.train_pipeline_config_path).apply_rollout_sampling_params( + args, sampling_params, extra_sampling_params + ) + if extra_sampling_params: sampling_params["extra_sampling_params"] = extra_sampling_params diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 5451ad46..a5b749c9 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -294,6 +294,24 @@ def add_rollout_arguments(parser): default=512, help="Width of the trained media: rollout output width, SFT encode center-crop width.", ) + parser.add_argument( + "--diffusion-h3-aspect-ratio", + type=str, + default="16:9", + help="MiniMax H3 target aspect_ratio: one of 21:9, 16:9, 4:3, 1:1, 3:4, 9:16.", + ) + parser.add_argument( + "--diffusion-h3-duration-seconds", + type=float, + default=5.0, + help="MiniMax H3 target duration_seconds for rollout (4.0-15.0).", + ) + parser.add_argument( + "--diffusion-audio-flow-shift", + type=float, + default=3.0, + help="MiniMax H3 audio_flow_shift for rollout.", + ) parser.add_argument( "--diffusion-negative-prompt", type=str, diff --git a/miles/utils/diffusion_rollout_response.py b/miles/utils/diffusion_rollout_response.py index 0163adf8..4b1eb6ba 100644 --- a/miles/utils/diffusion_rollout_response.py +++ b/miles/utils/diffusion_rollout_response.py @@ -107,12 +107,38 @@ def _parse_cond_kwargs( data.get("audio_encoder_attention_mask"), deserialize_func=deserialize_func ), pooled_projections=_parse_tensor_or_list(data.get("pooled_projections"), deserialize_func=deserialize_func), + h3_packed_layout=_parse_h3_packed_layout(data.get("h3_packed_layout"), deserialize_func=deserialize_func), + h3_token_tags=_deserialize_optional_tensor(data.get("h3_token_tags"), deserialize_func=deserialize_func), text_ids=deserialize_func(data.get("text_ids")), text_mask=deserialize_func(data.get("text_mask")), fps=data.get("fps"), ) +def _deserialize_optional_tensor(value, *, deserialize_func): + if value is None: + return None + if isinstance(value, dict) and value.get("__tensor__"): + return deserialize_func(value) + if isinstance(value, torch.Tensor): + return value + return value + + +def _parse_h3_packed_layout(value, *, deserialize_func): + if value is None: + return None + if not isinstance(value, dict): + return value + out = {} + for key, item in value.items(): + if isinstance(item, dict) and item.get("__tensor__"): + out[key] = deserialize_func(item) + else: + out[key] = item + return out + + def _parse_denoising_env( data: dict[str, Any] | None, *, diff --git a/miles/utils/types.py b/miles/utils/types.py index 2e6d900f..c4b972f6 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -33,6 +33,9 @@ class CondKwargs: encoder_attention_mask: list[torch.Tensor] | None = None audio_encoder_attention_mask: list[torch.Tensor] | None = None pooled_projections: list[torch.Tensor] | None = None + # MiniMax H3 packed-sequence replay metadata (from rollout denoising_env). + h3_packed_layout: dict | None = None + h3_token_tags: torch.Tensor | None = None # Cosmos3: token-level conditioning (no separate text encoder). text_ids: torch.Tensor | None = None text_mask: torch.Tensor | None = None diff --git a/miles/utils/wandb_utils.py b/miles/utils/wandb_utils.py index b9fc445c..9f7027bc 100644 --- a/miles/utils/wandb_utils.py +++ b/miles/utils/wandb_utils.py @@ -118,17 +118,6 @@ def init_wandb_secondary(args, router_addr=None): x_update_finish_state=False, ) - if getattr(args, "sglang_enable_metrics", False) and router_addr is not None: - logger.info(f"Forward SGLang metrics at {router_addr} to WandB.") - settings_kwargs |= dict( - x_stats_open_metrics_endpoints={ - "sgl_engine": f"{router_addr}/engine_metrics", - }, - x_stats_open_metrics_filters={ - "sgl_engine.*": {}, - }, - ) - init_kwargs = { "id": wandb_run_id, "entity": args.wandb_team, diff --git a/scripts/run_diffusion_grpo_h3_t2va_2gpu.py b/scripts/run_diffusion_grpo_h3_t2va_2gpu.py new file mode 100644 index 00000000..bb808913 --- /dev/null +++ b/scripts/run_diffusion_grpo_h3_t2va_2gpu.py @@ -0,0 +1,208 @@ +"""MiniMax H3 t2va PickScore Flow-GRPO: 2-GPU FSDP train + sglang rollout, reward colocated. + +Ported from the verl-omni DanceGRPO task18 recipe. Two of its settings cannot be +reproduced here and are deliberately different: + * Resolution/frames: sglang H3 pins short_edge=768 and duration >= 4s with 17n+5 + frame alignment at 24 fps, so the smallest legal spec is 1344x768 / 107 frames. + verl's 480x864 / 56f (2.33s) is not a valid H3 request. + * 2 GPUs instead of 8, so batch x group is smaller than verl's 8x16. + +The SDE window starts at step 1, not 0: the flow-SDE diffusion factor carries a +1/(1-sigma) that is singular at the first step (sigma=1), where the rollout engine +silently emits NaN latents. Excluding that one step is what lets this run use sde +(matching verl) instead of falling back to cps. + +Aligned with verl: LoRA 64/128, lr 1e-4, weight_decay 1e-4, adam_eps 1e-15, +noise_level 0.7, 10 inference steps, SDE window size 2, flow_shift 12.0 / +audio_flow_shift 3.0, PickScore over 8 frames. + +The router health-check window is widened well past its 30s x 3 default: H3 denoising +blocks the engine's uvicorn event loop for ~43s per sample, so /health cannot answer +while a request runs. A 16-sample eval keeps the loop busy for ~800s straight, which +trips the default threshold and quarantines a perfectly healthy engine permanently +(router.py marks DEAD with no revive path). + +Batch math: rollout_batch_size prompts x n_samples_per_prompt samples per rollout, +split into num_steps_per_rollout optimizer steps. The GRPO group is +n_samples_per_prompt, since advantages are standardized within a prompt's group, so a +small group leaves a weak, noisy signal. eval_interval and save_interval count rollout +iterations, not optimizer steps. + +Usage: + python3 scripts/run_diffusion_grpo_h3_t2va_2gpu.py + python3 scripts/run_diffusion_grpo_h3_t2va_2gpu.py --num-rollout 5 --cuda-visible-devices 0,3 + + # rollout-only smoke: sglang rollout + reward, no FSDP train + python3 scripts/run_diffusion_grpo_h3_t2va_2gpu.py --num-rollout 1 \ + --n-samples-per-prompt 1 --eval-interval 0 --extra-args "--debug-rollout-only" + + # train/rollout alignment diagnostic: freeze the weights so log_prob_mean_abs_diff + # measures pure train-vs-rollout deviation rather than parameter drift + python3 scripts/run_diffusion_grpo_h3_t2va_2gpu.py --num-rollout 2 \ + --eval-interval 0 --extra-args "--debug-skip-optimizer-step" +""" + +from dataclasses import dataclass +from pathlib import Path + +import typer + +import miles.utils.external_utils.command_utils as U + +MODEL = "MiniMaxAI/MiniMax-H3" +DATASET = "rockdu/miles-diffusion-datasets" +DATASET_SUBSET = "flowgrpo_pickscore" +WANDB_PROJECT = "miles-diffusion-grpo" + + +@dataclass +class ScriptArgs(U.ExecuteTrainConfig): + cuda_visible_devices: str = "0,1" + num_rollout: int = 30 + rollout_batch_size: int = 2 + n_samples_per_prompt: int = 8 + num_steps_per_rollout: int = 2 + eval_interval: int = 10 + # miles has no eval-sample cap (verl used val_max_samples=64), so a fixed slice of + # the test split stands in: all 2048 prompts would take >60h at 20 steps here. + eval_size: int = 16 + save_interval: int = 10 + data_dir: str = "/root/datasets" + extra_args: str = "" + + +def prepare(args: ScriptArgs) -> str: + local_dir = U.hf_download_dataset(DATASET, include=f"{DATASET_SUBSET}/**", data_dir=args.data_dir) + return f"{local_dir}/{DATASET_SUBSET}" + + +def _eval_slice(prompt_dir: str, eval_size: int) -> str: + eval_data = Path(prompt_dir) / f"val_{eval_size}.jsonl" + if not eval_data.exists(): + with open(Path(prompt_dir) / "test.jsonl") as f: + lines = [next(f) for _ in range(eval_size)] + eval_data.write_text("".join(lines)) + return str(eval_data) + + +def execute(args: ScriptArgs, prompt_dir: str) -> None: + run_name = f"diffusion_grpo_h3_t2va_{U.create_run_id()}" + + ckpt_args = ( + f"--hf-checkpoint {MODEL} " + f"--save {args.output_dir}/{run_name}/ckpt " + f"--save-interval {args.save_interval} " + ) + + rollout_args = ( + "--rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout " + f"--prompt-data {prompt_dir}/train.jsonl " + "--input-key input " + f"--rollout-batch-size {args.rollout_batch_size} " + f"--n-samples-per-prompt {args.n_samples_per_prompt} " + f"--num-steps-per-rollout {args.num_steps_per_rollout} " + f"--num-rollout {args.num_rollout} " + "--rollout-microgroup-size 1 " + "--micro-batch-size-sample 1 " + "--micro-batch-size-tstep 1 " + "--diffusion-train-iter-order sample_major " + ) + + diffusion_args = ( + "--diffusion-num-steps 10 " + "--diffusion-guidance-scale 1.0 " + "--diffusion-noise-level 0.7 " + "--diffusion-sde-type sde " + "--diffusion-step-strategy-path miles.rollout.step_strategy_hub.sde_window " + "--diffusion-num-sde-steps 2 " + "--diffusion-sde-window-range 1,4 " + "--diffusion-h3-aspect-ratio 16:9 " + "--diffusion-h3-duration-seconds 4 " + "--diffusion-audio-flow-shift 3.0 " + "--diffusion-flow-shift 12.0 " + ) + + eval_args = "" + if args.eval_interval: + eval_args = ( + f"--eval-interval {args.eval_interval} " + f"--eval-prompt-data pickscore_val {_eval_slice(prompt_dir, args.eval_size)} " + "--n-samples-per-eval-prompt 1 " + "--diffusion-eval-num-steps 20 " + ) + + grpo_args = ( + "--advantage-estimator grpo --globalize-reward-std --diffusion-clip-range 1e-4 --diffusion-kl-beta 0.0 " + ) + + optimizer_args = "--lr 1e-4 --weight-decay 1e-4 --adam-eps 1e-15 " + + # H3's rollout DiT renames modules and fuses Q/K/V, so weights only reach the engine + # through the LoRA IPC path's layer grouper; the family rejects any other sync mode. + lora_args = "--use-lora --lora-ipc-weight-sync --lora-rank 64 --lora-alpha 128 " + + reward_args = ( + "--rm-type pickscore " + "--pickscore-processor-path laion/CLIP-ViT-H-14-laion2B-s32B-b79K " + "--pickscore-model-path yuvalkirstain/PickScore_v1 " + "--pickscore-num-frames 8 " + "--pickscore-num-workers 1 " + "--pickscore-num-gpus-per-worker 0 " + "--pickscore-batch-size 4 " + "--rollout-parser-num-workers 2 " + ) + + wandb_args = U.get_default_wandb_args(__file__, run_id=run_name, project=WANDB_PROJECT, wandb_log_num_images=4) + + sglang_args = ( + "--use-miles-router " + "--sglang-server-concurrency 1 " + "--sglang-tp-size 2 " + "--sglang-sp-degree 1 " + "--sglang-ulysses-degree 1 " + "--sglang-ring-degree 1 " + "--sglang-dit-precision bf16 " + "--update-weight-buffer-size 2147483648 " + ) + + train_backend_args = ( + "--train-backend fsdp " + "--fsdp-master-dtype bf16 " + "--fsdp-reduce-dtype bf16 " + "--diffusion-forward-dtype bf16 " + ) + + perf_args = "--gradient-checkpointing " + + misc_args = ( + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 2 " + "--rollout-num-gpus 2 " + "--rollout-num-gpus-per-engine 2 " + "--num-gpus-per-node 2 " + "--colocate " + "--colocate-reward " + "--rollout-health-check-interval 60 " + "--miles-router-health-check-failure-threshold 30 " + ) + + U.execute_train( + train_args=( + f"{ckpt_args} {rollout_args} {diffusion_args} {eval_args} {grpo_args} {optimizer_args} " + f"{lora_args} {reward_args} {wandb_args} {sglang_args} {train_backend_args} {perf_args} " + f"{misc_args} {args.extra_args}" + ), + num_gpus_per_node=2, + config=args, + extra_env_vars={"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}, + ) + + +@U.dataclass_cli +def main(args: ScriptArgs) -> None: + prompt_dir = prepare(args) + execute(args, prompt_dir) + + +if __name__ == "__main__": + typer.run(main) diff --git a/tests/fast-gpu/test_lora_weight_sync.py b/tests/fast-gpu/test_lora_weight_sync.py index fb2d3e5f..192f4a26 100644 --- a/tests/fast-gpu/test_lora_weight_sync.py +++ b/tests/fast-gpu/test_lora_weight_sync.py @@ -59,14 +59,18 @@ def _make_peft_model(): return peft_model +def _fake_args(buffer_size): + return Namespace(update_weight_buffer_size=buffer_size, train_pipeline_config_path=None) + + def _run_update(peft_model, buffer_size): - updater = _CaptureUpdater(Namespace(update_weight_buffer_size=buffer_size), {"transformer": peft_model}) + updater = _CaptureUpdater(_fake_args(buffer_size), {"transformer": peft_model}) updater.update_weights() return updater.buckets def _run_lora_ipc_update(peft_model, buffer_size): - updater = _CaptureLoRAIPCUpdater(Namespace(update_weight_buffer_size=buffer_size), {"transformer": peft_model}) + updater = _CaptureLoRAIPCUpdater(_fake_args(buffer_size), {"transformer": peft_model}) updater.update_weights() return updater.buckets diff --git a/tests/fast/rollout/test_rollout_negative_prompt.py b/tests/fast/rollout/test_rollout_negative_prompt.py index 43bf95f5..265e52b3 100644 --- a/tests/fast/rollout/test_rollout_negative_prompt.py +++ b/tests/fast/rollout/test_rollout_negative_prompt.py @@ -30,6 +30,7 @@ def _args(**overrides): diffusion_noise_level=0.7, diffusion_log_prob_no_const=False, diffusion_debug_mode=False, + train_pipeline_config_path=None, ) values.update(overrides) return Namespace(**values) diff --git a/tests/fast/utils/test_lora_args.py b/tests/fast/utils/test_lora_args.py index c0434ff4..a3615f77 100644 --- a/tests/fast/utils/test_lora_args.py +++ b/tests/fast/utils/test_lora_args.py @@ -12,6 +12,7 @@ def _server_args(**overrides): hf_checkpoint="Qwen/Qwen-Image", diffusion_flow_shift=None, rollout_num_gpus_per_engine=1, + sglang_tp_size=1, sglang_sp_degree=None, sglang_enable_cfg_parallel=False, use_lora=True,