| """RLSD warmup gates for OPSD degenerate skip, denser online SFT, and embedded SFT cold start.""" |
| from __future__ import annotations |
|
|
| import math |
| from typing import Any, Mapping, Optional |
|
|
|
|
| def current_global_step(trainer: Any) -> int: |
| return int(getattr(getattr(trainer, "state", None), "global_step", getattr(trainer, "_step", 0)) or 0) |
|
|
|
|
| def resolve_max_training_steps(trainer: Any) -> Optional[int]: |
| """Resolve total optimizer steps for gate math (cold start frac, warmup windows). |
| |
| Priority: TrainingArguments.max_steps > Trainer.state.max_steps > epoch estimate. |
| HF sets state.max_steps when max_steps<=0 from num_train_epochs * len(dataloader). |
| """ |
| args = getattr(trainer, "args", None) |
| if args is not None: |
| arg_max = getattr(args, "max_steps", None) |
| if arg_max is not None and int(arg_max) > 0: |
| return int(arg_max) |
|
|
| state = getattr(trainer, "state", None) |
| if state is not None: |
| state_max = getattr(state, "max_steps", None) |
| if state_max is not None and int(state_max) > 0: |
| return int(state_max) |
|
|
| if args is None: |
| return None |
|
|
| num_epochs = getattr(args, "num_train_epochs", None) |
| grad_accum = max(1, int(getattr(args, "gradient_accumulation_steps", 1) or 1)) |
| if num_epochs is None or float(num_epochs) <= 0: |
| return None |
|
|
| dataloader = getattr(trainer, "train_dataloader", None) |
| if dataloader is None and hasattr(trainer, "get_train_dataloader"): |
| try: |
| dataloader = trainer.get_train_dataloader() |
| except Exception: |
| dataloader = None |
| if dataloader is None: |
| return None |
|
|
| try: |
| steps_per_epoch = len(dataloader) |
| except TypeError: |
| return None |
| if steps_per_epoch <= 0: |
| return None |
|
|
| total = math.ceil(float(num_epochs) * steps_per_epoch / grad_accum) |
| return total if total > 0 else None |
|
|
|
|
| def sft_cold_start_steps(opsd_config: Mapping[str, Any], max_steps: Optional[int]) -> int: |
| """Steps at start of training devoted to embedded offline-style SFT (no generate / no OPSD).""" |
| gate = opsd_config.get("gate", {}) |
| steps_env = gate.get("sft_cold_start_steps") |
| if steps_env is not None: |
| return max(0, int(steps_env)) |
| frac = float(gate.get("sft_cold_start_frac", 0.0) or 0.0) |
| if frac <= 0.0 or max_steps is None or max_steps <= 0: |
| return 0 |
| return max(1, int(max_steps * frac)) |
|
|
|
|
| def in_sft_cold_start( |
| opsd_config: Mapping[str, Any], |
| global_step: int, |
| max_steps: Optional[int], |
| ) -> bool: |
| cold_steps = sft_cold_start_steps(opsd_config, max_steps) |
| return cold_steps > 0 and global_step < cold_steps |
|
|
|
|
| def resolve_skip_degenerate_opsd( |
| opsd_config: Mapping[str, Any], |
| global_step: int, |
| max_steps: Optional[int] = None, |
| ) -> bool: |
| gate = opsd_config.get("gate", {}) |
| if not gate.get("skip_degenerate_for_opsd", False): |
| return False |
| cold_end = sft_cold_start_steps(opsd_config, max_steps) |
| warmup = int(gate.get("degen_skip_warmup_steps", 200)) |
| |
| threshold = cold_end + warmup if cold_end > 0 else warmup |
| return global_step >= threshold |
|
|
|
|
| def sft_slots_for_step( |
| opsd_config: Mapping[str, Any], |
| global_step: int, |
| max_steps: Optional[int] = None, |
| ) -> int: |
| if in_sft_cold_start(opsd_config, global_step, max_steps): |
| return 0 |
| gate = opsd_config.get("gate", {}) |
| warmup_steps = int(gate.get("sft_warmup_steps", 200)) |
| cold_end = sft_cold_start_steps(opsd_config, max_steps) |
| effective_warmup_end = cold_end + warmup_steps if cold_end > 0 else warmup_steps |
| if global_step < effective_warmup_end: |
| return max(1, int(gate.get("sft_warmup_slots_per_group", 2))) |
| return 1 |
|
|
