From ace421bc2bdb375dda20ff2639785770328b6ed2 Mon Sep 17 00:00:00 2001 From: Mac Lee Date: Fri, 19 Jun 2026 16:45:33 +0000 Subject: [PATCH] [feat] harden InterleaveThinker critic backend --- ...interleavethinker-fastvideo-integration.md | 92 +++++++++++ .../rl/interleave_thinker/critic_grpo.yaml | 1 + .../train/models/interleave_thinker/critic.py | 85 ++++++++-- .../test_interleave_thinker_critic_model.py | 153 +++++++++++++++++- .../test_interleave_thinker_method.py | 1 + 5 files changed, 317 insertions(+), 15 deletions(-) diff --git a/.agents/exploration/interleavethinker-fastvideo-integration.md b/.agents/exploration/interleavethinker-fastvideo-integration.md index 01ff9c88d..f5817fb30 100644 --- a/.agents/exploration/interleavethinker-fastvideo-integration.md +++ b/.agents/exploration/interleavethinker-fastvideo-integration.md @@ -435,6 +435,98 @@ Pending: - Commit this model/API integration slice. - Push the commit immediately after committing, per user instruction. +Commit completed: + +- `fc04d019` — `[feat] integrate InterleaveThinker model backends` +- Pushed to `origin/interleavethinker-fastvideo`. + +## Critic Backend Hardening Extension + +Status: implemented and Modal-validated; pending commit/push. + +Goal for this slice: + +- Tighten the Qwen3-VL critic wrapper from config/import-ready toward + executable-backend-ready without pulling real Qwen weights into unit tests. +- Preserve user constraints: + - work only in `/tmp/fastvideo-interleavethinker`, + - do not touch `/home/toolbox/FastVideo`, + - validate via Modal, + - commit and push immediately after commit. + +Planned changes: + +1. Add `image_dir` handling to `InterleaveThinkerCriticModel`, matching + upstream EasyR1's convention that relative `origin_image_path` and + `edited_image_path` are resolved under a configured image root. +2. Normalize image paths at dataset load/collate time, and expose an explicit + before/after image pair: + - before: `previous_image_path`, then `origin_image_path`, then + `input_image_path`; + - after: `edited_image_path`, then `generated_image_path`, then + `output_image_path`. +3. Improve actor-update scaling so per-rollout backward computes the mean batch + policy gradient instead of summing all rollout losses. +4. Add fake Qwen backend tests that exercise: + - `generate_interleave_responses(...)`, + - `train_interleave_rollouts(...)`, + - image path resolution, + - public YAML parse for the new config knobs. + +Validation plan: + +- Modal pytest: + `pytest tests/local_tests/test_interleave_thinker_reward.py + tests/local_tests/test_interleave_thinker_method.py + tests/local_tests/test_interleave_thinker_api_models.py + tests/local_tests/test_interleave_thinker_critic_model.py + tests/local_tests/test_train_rl_sampling.py -q` +- Modal pre-commit on changed files. + +Implemented changes: + +- `InterleaveThinkerCriticModel` accepts `image_dir` and passes it into the + JSON/JSONL dataset loader. +- Dataset loading normalizes relative image paths under `image_dir` while + leaving absolute paths, URI-style paths, and `data:` URLs unchanged. +- Qwen3-VL message assembly now selects an explicit before/after image pair. + This avoids sending `origin`, `previous`, and `edited` together when the + critic task expects two images. +- Actor update scaling now averages over `len(rollouts) * + gradient_accumulation_steps`, preventing the loss from growing with rollout + count. +- Added fake-backend tests for generation and response-token-only actor + training without downloading real Qwen weights. +- Added `image_dir: data` to + `examples/train/configs/rl/interleave_thinker/critic_grpo.yaml`. + +Validation completed: + +- Modal pytest: + - App URL: `https://modal.com/apps/hao-ai-lab/main/ap-fRlUUdFx19OWpSvDSrWoAe` + - Command: + `pytest tests/local_tests/test_interleave_thinker_reward.py + tests/local_tests/test_interleave_thinker_method.py + tests/local_tests/test_interleave_thinker_api_models.py + tests/local_tests/test_interleave_thinker_critic_model.py + tests/local_tests/test_train_rl_sampling.py -q` + - Result: `30 passed, 14 warnings in 19.47s`. +- Modal pre-commit: + - App URL: `https://modal.com/apps/hao-ai-lab/main/ap-DplMFq23YYfBx34e6TcsRc` + - Command: + `pre-commit run --files + examples/train/configs/rl/interleave_thinker/critic_grpo.yaml + fastvideo/train/models/interleave_thinker/critic.py + tests/local_tests/test_interleave_thinker_critic_model.py + tests/local_tests/test_interleave_thinker_method.py` + - Result: yapf, ruff, codespell, mypy, filename check, and suggestion hooks + passed; PyMarkdown/actionlint skipped with no files to check. + +Pending: + +- Commit this critic-backend hardening slice. +- Push the commit immediately after committing, per user instruction. + ## Training Loop Extension Status: in progress as of the user request to "Add a training loop" for diff --git a/examples/train/configs/rl/interleave_thinker/critic_grpo.yaml b/examples/train/configs/rl/interleave_thinker/critic_grpo.yaml index 5c924c252..ebb84ff3d 100644 --- a/examples/train/configs/rl/interleave_thinker/critic_grpo.yaml +++ b/examples/train/configs/rl/interleave_thinker/critic_grpo.yaml @@ -11,6 +11,7 @@ models: init_from: InterleaveThinker/Critic-SFT-8B processor_from: Qwen/Qwen3-VL-8B-Instruct load_backend: true + image_dir: data torch_dtype: auto attn_implementation: flash_attention_2 freeze_vision_tower: true diff --git a/fastvideo/train/models/interleave_thinker/critic.py b/fastvideo/train/models/interleave_thinker/critic.py index 1356dff36..05b236fa9 100644 --- a/fastvideo/train/models/interleave_thinker/critic.py +++ b/fastvideo/train/models/interleave_thinker/critic.py @@ -60,8 +60,13 @@ class _InterleaveJSONLDataset(Dataset): def __init__( self, data_path: str, + *, + image_dir: str = "", ) -> None: - self.records = list(_load_interleave_records(data_path)) + self.records = list(_load_interleave_records( + data_path, + image_dir=image_dir, + )) if not self.records: raise ValueError(f"No InterleaveThinker records found at {data_path!r}") @@ -101,6 +106,7 @@ class InterleaveThinkerCriticModel(ModelBase): training_config: TrainingConfig | None = None, trainable: bool = True, load_backend: bool = True, + image_dir: str = "", torch_dtype: str = "auto", device_map: str | dict[str, Any] | None = None, attn_implementation: str | None = None, @@ -119,6 +125,7 @@ class InterleaveThinkerCriticModel(ModelBase): super().__init__(trainable=trainable, lora=lora) self.init_from = str(init_from) self.processor_from = str(processor_from or init_from) + self.image_dir = str(image_dir or "") self.training_config = training_config self.max_prompt_length = int(max_prompt_length) self.max_response_length = int(max_response_length) @@ -214,7 +221,10 @@ class InterleaveThinkerCriticModel(ModelBase): return if not os.path.exists(os.path.expanduser(data_path)): raise FileNotFoundError(f"InterleaveThinker data_path not found: {data_path}") - dataset = _InterleaveJSONLDataset(data_path) + dataset = _InterleaveJSONLDataset( + data_path, + image_dir=self.image_dir, + ) self.dataloader = DataLoader( dataset, batch_size=max(1, int(training_config.data.train_batch_size or 1)), @@ -295,10 +305,11 @@ class InterleaveThinkerCriticModel(ModelBase): weighted_losses: list[torch.Tensor] = [] nll_values: list[torch.Tensor] = [] token_counts: list[int] = [] + loss_scale = max(1, len(rollouts) * grad_accum) for idx, rollout in enumerate(rollouts): nll, token_count = self._response_nll(rollout) weighted = nll * advantages[idx] - (weighted / grad_accum).backward() + (weighted / loss_scale).backward() weighted_losses.append(weighted.detach()) nll_values.append(nll.detach()) token_counts.append(token_count) @@ -446,7 +457,11 @@ class InterleaveThinkerCriticModel(ModelBase): raise NotImplementedError("InterleaveThinkerCriticModel uses train_interleave_rollouts().") -def _load_interleave_records(data_path: str, ) -> list[dict[str, Any]]: +def _load_interleave_records( + data_path: str, + *, + image_dir: str = "", +) -> list[dict[str, Any]]: path = Path(os.path.expanduser(data_path)) files = sorted(path.glob("*.json*")) if path.is_dir() else [path] records: list[dict[str, Any]] = [] @@ -456,18 +471,55 @@ def _load_interleave_records(data_path: str, ) -> list[dict[str, Any]]: if line.strip(): raw = json.loads(line) if isinstance(raw, Mapping): - records.append(dict(raw)) + records.append(_normalize_interleave_record(raw, image_dir=image_dir)) elif file_path.suffix == ".json": raw_json = json.loads(file_path.read_text()) if isinstance(raw_json, list): - records.extend(dict(item) for item in raw_json if isinstance(item, Mapping)) + records.extend( + _normalize_interleave_record(item, image_dir=image_dir) for item in raw_json + if isinstance(item, Mapping)) elif isinstance(raw_json, Mapping): - records.append(dict(raw_json)) + records.append(_normalize_interleave_record(raw_json, image_dir=image_dir)) else: raise ValueError(f"Unsupported InterleaveThinker data file: {file_path}") return records +def _normalize_interleave_record( + record: Mapping[str, Any], + *, + image_dir: str = "", +) -> dict[str, Any]: + normalized = dict(record) + for key in ( + "origin_image_path", + "previous_image_path", + "edited_image_path", + "generated_image_path", + "input_image_path", + "output_image_path", + ): + value = normalized.get(key) + if isinstance(value, str) and value: + normalized[key] = _resolve_image_path(value, image_dir=image_dir) + return normalized + + +def _resolve_image_path( + value: str, + *, + image_dir: str = "", +) -> str: + expanded = os.path.expanduser(value) + if not image_dir or os.path.isabs(expanded) or _looks_like_uri(expanded): + return expanded + return str(Path(os.path.expanduser(image_dir)) / expanded) + + +def _looks_like_uri(value: str) -> bool: + return "://" in value or value.startswith("data:") + + def _collate_interleave_records(records: Sequence[Mapping[str, Any]], ) -> dict[str, Any]: return {"items": [dict(record) for record in records]} @@ -493,16 +545,25 @@ def _batch_to_items(batch: Mapping[str, Any], ) -> list[dict[str, Any]]: def _item_image_paths(item: Mapping[str, Any], ) -> list[str]: - image_paths: list[str] = [] - for key in ("origin_image_path", "previous_image_path", "edited_image_path"): - value = item.get(key) - if isinstance(value, str) and value: - image_paths.append(value) + before = _first_string(item, "previous_image_path", "origin_image_path", "input_image_path") + after = _first_string(item, "edited_image_path", "generated_image_path", "output_image_path") + image_paths = [value for value in (before, after) if value] if len(image_paths) == 1: image_paths.append(image_paths[0]) return image_paths[:2] +def _first_string( + item: Mapping[str, Any], + *keys: str, +) -> str | None: + for key in keys: + value = item.get(key) + if isinstance(value, str) and value: + return value + return None + + def _move_to_device( value: Any, device: torch.device, diff --git a/tests/local_tests/test_interleave_thinker_critic_model.py b/tests/local_tests/test_interleave_thinker_critic_model.py index bec6a091b..0f818b632 100644 --- a/tests/local_tests/test_interleave_thinker_critic_model.py +++ b/tests/local_tests/test_interleave_thinker_critic_model.py @@ -1,5 +1,11 @@ import json +from types import SimpleNamespace +import torch + +from fastvideo.train.models.interleave_thinker.critic import ( + _PlaceholderActorModule, +) from fastvideo.train.models.interleave_thinker import ( INTERLEAVE_CRITIC_PROMPT, InterleaveThinkerCriticModel, @@ -10,6 +16,65 @@ from fastvideo.train.utils.training_config import ( ) +class _FakeBackendCritic(InterleaveThinkerCriticModel): + + @property + def device(self): + return torch.device("cpu") + + +class _FakeProcessor: + + def __init__(self): + self.tokenizer = SimpleNamespace(pad_token_id=0, eos_token_id=2) + + def apply_chat_template( + self, + messages, + *, + tokenize, + add_generation_prompt, + return_dict, + return_tensors, + ): + del tokenize, add_generation_prompt, return_dict, return_tensors + has_assistant = any(message["role"] == "assistant" for message in messages) + length = 5 if has_assistant else 3 + return { + "input_ids": torch.arange(1, length + 1).unsqueeze(0), + "attention_mask": torch.ones(1, length, dtype=torch.long), + } + + def batch_decode(self, sequences, **kwargs): + del kwargs + return [ + 'ok{"previous_step_success": true, "refine_prompt": "better"}' + for _ in sequences + ] + + +class _FakeQwen(torch.nn.Module): + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor(1.0)) + self.generate_kwargs = None + self.last_labels = None + + def generate(self, **kwargs): + self.generate_kwargs = kwargs + input_ids = kwargs["input_ids"] + num_return_sequences = int(kwargs["num_return_sequences"]) + suffix = torch.tensor([[9, 10]], dtype=input_ids.dtype) + return torch.cat([input_ids.repeat(num_return_sequences, 1), suffix.repeat(num_return_sequences, 1)], dim=1) + + def forward(self, **kwargs): + labels = kwargs["labels"] + self.last_labels = labels.detach().clone() + trainable_fraction = (labels != -100).float().mean() + return SimpleNamespace(loss=self.weight.pow(2).sum() * trainable_fraction) + + def test_interleave_thinker_critic_builds_qwen_vl_messages_without_loading_backend(): model = InterleaveThinkerCriticModel(load_backend=False) @@ -35,25 +100,107 @@ def test_interleave_thinker_critic_builds_qwen_vl_messages_without_loading_backe assert "a vase on a table" in text +def test_interleave_thinker_critic_prefers_previous_image_as_before_image(): + model = InterleaveThinkerCriticModel(load_backend=False) + + messages = model.build_messages({ + "origin_prompt": "draw a vase", + "previous_prompt": "a vase on a table", + "origin_image_path": "original.png", + "previous_image_path": "before.png", + "edited_image_path": "after.png", + }) + + images = [part for part in messages[0]["content"] if part["type"] == "image"] + assert images == [{ + "type": "image", + "image": "before.png", + }, { + "type": "image", + "image": "after.png", + }] + + def test_interleave_thinker_critic_initializes_jsonl_dataloader(tmp_path): data_path = tmp_path / "critic_rl.jsonl" data_path.write_text( json.dumps({ "origin_prompt": "draw a chair", "rewritten_prompt": "a wooden chair", - "origin_image_path": "before.png", - "edited_image_path": "after.png", + "origin_image_path": "interleave/before.png", + "edited_image_path": "interleave/after.png", "evaluation": { "success": True, "semantics": 7.0, "quality": 8.0, }, }) + "\n") - model = InterleaveThinkerCriticModel(load_backend=False) + image_dir = tmp_path / "images" + model = InterleaveThinkerCriticModel(load_backend=False, image_dir=str(image_dir)) model.init_preprocessors(TrainingConfig(data=DataConfig(data_path=str(data_path), train_batch_size=1))) batch = next(iter(model.dataloader)) assert batch["items"][0]["origin_prompt"] == "draw a chair" + assert batch["items"][0]["origin_image_path"] == str(image_dir / "interleave/before.png") + assert batch["items"][0]["edited_image_path"] == str(image_dir / "interleave/after.png") assert batch["items"][0]["evaluation"]["success"] is True assert "{original_instruction}" in INTERLEAVE_CRITIC_PROMPT + + +def test_interleave_thinker_critic_fake_backend_generates_rollouts(): + model = _FakeBackendCritic(load_backend=False) + model.processor = _FakeProcessor() + model.transformer = _FakeQwen() + + rollouts = model.generate_interleave_responses( + { + "items": [{ + "origin_prompt": "draw a chair", + "previous_prompt": "a wooden chair", + "origin_image_path": "before.png", + "edited_image_path": "after.png", + }] + }, + num_generations=2, + temperature=0.7, + top_p=0.9, + max_new_tokens=8, + ) + + assert len(rollouts) == 2 + assert all("previous_step_success" in rollout["response"] for rollout in rollouts) + assert rollouts[0]["group_key"] == rollouts[1]["group_key"] + assert rollouts[0]["sample_index"] == 0 + assert model.transformer.generate_kwargs["num_return_sequences"] == 2 + assert model.transformer.generate_kwargs["max_new_tokens"] == 8 + + +def test_interleave_thinker_critic_fake_backend_trains_response_tokens_only(): + model = _FakeBackendCritic(load_backend=False) + assert isinstance(model.transformer, _PlaceholderActorModule) + model.processor = _FakeProcessor() + model.transformer = _FakeQwen() + optimizer = torch.optim.SGD(model.transformer.parameters(), lr=0.1) + before = float(model.transformer.weight.detach()) + + loss_map, metrics = model.train_interleave_rollouts( + rollouts=[{ + "origin_prompt": "draw a chair", + "previous_prompt": "a wooden chair", + "origin_image_path": "before.png", + "edited_image_path": "after.png", + "response": 'ok{"previous_step_success": true, "refine_prompt": "better"}', + }], + advantages=torch.tensor([1.0]), + rewards={}, + optimizer=optimizer, + gradient_accumulation_steps=1, + max_grad_norm=0.0, + ) + + assert "total_loss" in loss_map + assert metrics["actor/response_tokens"] == 2.0 + assert float(model.transformer.weight.detach()) < before + assert model.transformer.last_labels.tolist()[0][:3] == [-100, -100, -100] + assert all(label != -100 for label in model.transformer.last_labels.tolist()[0][3:]) diff --git a/tests/local_tests/test_interleave_thinker_method.py b/tests/local_tests/test_interleave_thinker_method.py index c46b24c1c..992f70bdb 100644 --- a/tests/local_tests/test_interleave_thinker_method.py +++ b/tests/local_tests/test_interleave_thinker_method.py @@ -218,6 +218,7 @@ def test_interleave_thinker_config_parses_public_yaml(): assert cfg.models["student"]["_target_"] == ( "fastvideo.train.models.interleave_thinker.InterleaveThinkerCriticModel") assert cfg.models["student"]["init_from"] == "InterleaveThinker/Critic-SFT-8B" + assert cfg.models["student"]["image_dir"] == "data" assert cfg.method["_target_"] == "fastvideo.train.methods.rl.interleave_thinker.InterleaveThinkerRLMethod" assert cfg.method["edit_scorer"]["_target_"] == ( "fastvideo.train.methods.rl.rewards.GeminiNanoBananaEditScorer")