SolitaryThinker
3ec511419f
Stress-tests the v2 Card/Loop/Program design with a UniRL/PromptRL-style joint RL recipe: a prompt-refiner LM expert and a flow generator expert, two separate experts driven by two loop types in one request, both updated simultaneously from a single RL reward. - loop/sampler.py: flow_sde_step_with_logprob — FlowGRPO SDE rollout sampler (per-step Gaussian log-prob), distinct from the deterministic ODE serve step. - request/params.py: gated sde_rollout/sde_noise_scale on DiffusionParams so the serve path stays byte-identical (default ODE). - models/wan21/loop.py: gated SDE-rollout capture in WanDenoiseLoop.advance. - models/backend.py: ToyPromptRefiner — a real REINFORCE categorical policy (the Qwen role), separate weights from the generator. - models/unified/: the unified card+program — two disjoint experts (llm + transformer) on ar_decode + diffusion_denoise; the topological opposite of the Cosmos3 MoT card, same vocabulary. - training/methods/unified_rl.py: joint GRPO — one reward -> group advantage -> LM token policy gradient + DiT FlowGRPO PPO; two LRs; prompt-only/joint flag; reuses the shared diffusion loop for rollout. - training/weight_sync.py: WeightSyncPlan gains a component scope so the two experts version + cache-invalidate independently (LM sync never flushes the frozen text-encoder feature cache). - tests/test_unified_rl.py: 9 tests incl. likelihood-based C2 identity, the two-loop interleave parity gate, joint vs prompt-only. Full suite 91 passed. designv4.md: unified design doc reflecting v2 as built+tested, with the joint RL stress test as the validating case study (the design held — new card + new method, no new runtime primitive).
59 lines
3.0 KiB
Python
59 lines
3.0 KiB
Python
"""WeightSyncPlan — versioned weight transfer with a declared role (design_v3 §10).
|
|
|
|
> RL ships a *role*, not "the weights": student / EMA / decay-blended old policy is declared
|
|
> (the landed NFT behavior policy is the *old* copy, not the student — the plan must carry that).
|
|
|
|
The lifecycle (design_v3 §10): freeze admission → drain/boundary-stop → transfer → bump
|
|
``weights_version`` → invalidate incompatible caches/graphs → publish → resume. Computed from
|
|
three inputs (mesh specs + per-model layout adapters + transport), validated pre-flight,
|
|
CPU-testable on fake pools. Here the "transfer" copies/blends the toy DiT params.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import itertools
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
|
|
class WeightRole(str, Enum):
|
|
STUDENT = "student" # the current learnable weights
|
|
EMA = "ema" # exponential moving average of student
|
|
OLD_POLICY = "old_policy" # decay-blended behavior policy (NFT samples from THIS, not student)
|
|
REFERENCE = "reference" # frozen KL anchor
|
|
TEACHER = "teacher" # frozen distillation teacher
|
|
CRITIC = "critic" # trainable fake-score / value head
|
|
|
|
|
|
_ver = itertools.count(1)
|
|
|
|
|
|
@dataclass
|
|
class WeightSyncPlan:
|
|
role: WeightRole
|
|
# three inputs (design_v3 §10): mesh specs (trainer↔engine), layout adapters, transport
|
|
mesh_specs: dict[str, Any] = field(default_factory=dict)
|
|
layout_adapters: dict[str, Any] = field(default_factory=dict)
|
|
transport: str = "in_proc" # in_proc | shm | nccl | rdma
|
|
decay: float = 0.0 # for EMA / OLD_POLICY blend
|
|
# which component(s) this plan syncs — versioned + cache-invalidated in isolation (design_v3 §7.1).
|
|
# Joint RL (UniRL) ships ONE plan per trainable expert: ``("transformer",)`` and ``("llm",)`` —
|
|
# so the LM weight sync does NOT flush the frozen text-encoder's feature cache, and vice-versa.
|
|
components: tuple[str, ...] = ("transformer",)
|
|
|
|
def apply(self, src_dit: Any, dst_dit: Any, dst_instance: Any = None) -> str:
|
|
"""Execute the lifecycle. Returns the new weights_version published on dst_instance."""
|
|
# 1) freeze admission + 2) drain/boundary-stop are scheduler ops (no-op in single-thread mini)
|
|
# 3) transfer per role
|
|
if self.role in (WeightRole.EMA, WeightRole.OLD_POLICY):
|
|
dst_dit.blend_from(src_dit, self.decay) # decay·dst + (1-decay)·src
|
|
else:
|
|
dst_dit.copy_from(src_dit) # hard copy (student push / reference init)
|
|
# 4) bump version + 5) invalidate ONLY this plan's component caches/graphs + 6) publish.
|
|
# Scoped to ``self.components`` so frozen text/vision encoders keep their feature caches,
|
|
# and the two experts of a joint-RL update version independently (design_v3 §7.1, §10).
|
|
version = f"w{next(_ver)}"
|
|
if dst_instance is not None:
|
|
dst_instance.set_weights_version(version, components=list(self.components))
|
|
return version
|