streaming_long_tuning ¶

def parse_multi_phased_distill_schedule(
    raw: Any,
    *,
    default_num_latent_t: int,
    default_streaming_chunk_size: int | None = None,
    default_streaming_max_length: int | None = None,
) -> list[DistillStage]:
    """Parse a compact/list multi-stage schedule.

    Preferred YAML form:

    ``[{stage: self_forcing, end_step: 700, num_latent_t: 21}, ...]``

    Compact string form accepts ``"700:21,3000:240"`` and treats the
    first stage as non-streaming self-forcing and later stages as streaming.
    """

    if raw is None or raw == "":
        max_length = default_streaming_max_length or default_num_latent_t
        return [
            DistillStage(
                name="streaming_long",
                start_step=0,
                end_step=None,
                num_latent_t=int(max_length),
                streaming_training=True,
                streaming_chunk_size=default_streaming_chunk_size,
                streaming_max_length=int(max_length),
            )
        ]

    stages: list[DistillStage] = []
    previous_end = 0

    if isinstance(raw, str):
        parts = [part.strip() for part in raw.split(",") if part.strip()]
        for idx, part in enumerate(parts):
            fields = [field.strip() for field in part.split(":")]
            end_step: int | None
            if len(fields) == 2:
                start_step = previous_end
                end_step = _as_int(
                    fields[0],
                    where="multi_phased_distill_schedule.end_step",
                )
                num_latent_t = _as_int(
                    fields[1],
                    where="multi_phased_distill_schedule.num_latent_t",
                )
            elif len(fields) == 3:
                start_step = _as_int(
                    fields[0],
                    where="multi_phased_distill_schedule.start_step",
                )
                end_step = _as_int(
                    fields[1],
                    where="multi_phased_distill_schedule.end_step",
                )
                num_latent_t = _as_int(
                    fields[2],
                    where="multi_phased_distill_schedule.num_latent_t",
                )
            else:
                raise ValueError("multi_phased_distill_schedule string entries must be "
                                 "'end_step:num_latent_t' or 'start:end:num_latent_t', "
                                 f"got {part!r}")
            streaming = idx > 0
            stages.append(
                DistillStage(
                    name="streaming_long" if streaming else "self_forcing",
                    start_step=start_step,
                    end_step=end_step,
                    num_latent_t=num_latent_t,
                    streaming_training=streaming,
                    streaming_chunk_size=(default_streaming_chunk_size if streaming else None),
                    streaming_max_length=num_latent_t if streaming else None,
                ))
            previous_end = end_step
    elif isinstance(raw, list | tuple):
        for idx, stage_raw in enumerate(raw):
            if not isinstance(stage_raw, dict):
                raise ValueError("multi_phased_distill_schedule list entries must be dicts")
            name = str(stage_raw.get("stage", "") or stage_raw.get("name", "")).strip()
            streaming_raw = stage_raw.get("streaming_training", None)
            if streaming_raw is None:
                streaming = name in {"streaming_long", "long", "streaming"}
            else:
                streaming = _as_bool(
                    streaming_raw,
                    where=f"multi_phased_distill_schedule[{idx}].streaming_training",
                )
            if not name:
                name = "streaming_long" if streaming else "self_forcing"

            start_step = _as_int(
                stage_raw.get("start_step", previous_end),
                where=f"multi_phased_distill_schedule[{idx}].start_step",
            )
            end_raw = stage_raw.get("end_step", None)
            end_step = (None if end_raw is None else _as_int(
                end_raw,
                where=f"multi_phased_distill_schedule[{idx}].end_step",
            ))
            num_latent_t = _as_int(
                stage_raw.get(
                    "num_latent_t",
                    stage_raw.get(
                        "streaming_max_length",
                        stage_raw.get("max_length", default_num_latent_t),
                    ),
                ),
                where=f"multi_phased_distill_schedule[{idx}].num_latent_t",
            )
            streaming_max = stage_raw.get("streaming_max_length", None)
            chunk_size = stage_raw.get("streaming_chunk_size", None)
            min_new = stage_raw.get("streaming_min_new_frame", None)
            fixed_overlap = stage_raw.get("streaming_fixed_overlap_latents", None)
            stages.append(
                DistillStage(
                    name=name,
                    start_step=start_step,
                    end_step=end_step,
                    num_latent_t=num_latent_t,
                    streaming_training=streaming,
                    streaming_chunk_size=(None if chunk_size is None else _as_int(
                        chunk_size,
                        where=("multi_phased_distill_schedule"
                               f"[{idx}].streaming_chunk_size"),
                    )),
                    streaming_max_length=(None if streaming_max is None else _as_int(
                        streaming_max,
                        where=("multi_phased_distill_schedule"
                               f"[{idx}].streaming_max_length"),
                    )),
                    streaming_min_new_frame=(None if min_new is None else _as_int(
                        min_new,
                        where=("multi_phased_distill_schedule"
                               f"[{idx}].streaming_min_new_frame"),
                    )),
                    streaming_fixed_overlap_latents=(None if fixed_overlap is None else _as_int(
                        fixed_overlap,
                        where=("multi_phased_distill_schedule"
                               f"[{idx}].streaming_fixed_overlap_latents"),
                    )),
                ))
            if end_step is not None:
                previous_end = end_step
    else:
        raise ValueError("multi_phased_distill_schedule must be a list, string, or empty")

    if not stages:
        raise ValueError("multi_phased_distill_schedule produced no stages")

    previous_end = 0
    for stage in stages:
        if stage.start_step < previous_end:
            raise ValueError("Distillation stages must be ordered and non-overlapping")
        if stage.end_step is not None and stage.end_step <= stage.start_step:
            raise ValueError("Distillation stage end_step must be > start_step")
        if stage.num_latent_t <= 0:
            raise ValueError("Distillation stage num_latent_t must be positive")
        if stage.streaming_training:
            max_length = stage.streaming_max_length or stage.num_latent_t
            chunk_size = stage.streaming_chunk_size or default_streaming_chunk_size
            if max_length <= 0:
                raise ValueError("streaming_max_length must be positive")
            if chunk_size is None or chunk_size <= 0:
                raise ValueError("streaming_chunk_size must be positive")
            if stage.streaming_fixed_overlap_latents is not None:
                if stage.streaming_fixed_overlap_latents < 0:
                    raise ValueError("streaming_fixed_overlap_latents must be non-negative")
                if stage.streaming_fixed_overlap_latents >= chunk_size:
                    raise ValueError("streaming_fixed_overlap_latents must be smaller than "
                                     "streaming_chunk_size")
        if stage.end_step is not None:
            previous_end = stage.end_step
    return stages

def select_distill_stage(
    stages: list[DistillStage],
    iteration: int,
) -> DistillStage:
    """Select the active stage for ``iteration``."""

    step = int(iteration)
    for stage in stages:
        if stage.end_step is None:
            if step >= stage.start_step:
                return stage
        elif stage.start_step <= step < stage.end_step:
            return stage
    return stages[-1]