Buckets:
| """FP32 EMA parameter view shared by supervised and on-policy trainers.""" | |
| from __future__ import annotations | |
| import contextlib | |
| from collections import defaultdict | |
| import torch | |
| class EMATeacher: | |
| """A lagged model view updated only after successful optimizer steps.""" | |
| STATE_SCHEMA = "shared_ema_teacher_v1" | |
| LEGACY_STATE_SCHEMAS = frozenset({"sft_ema_teacher_v1"}) | |
| def __init__(self, student, decay: float): | |
| if not 0.0 <= decay < 1.0: | |
| raise ValueError("EMA decay must be in [0, 1)") | |
| self.decay = float(decay) | |
| named = tuple( | |
| (name, parameter) | |
| for name, parameter in student.named_parameters() | |
| if parameter.requires_grad | |
| ) | |
| self.parameter_names = tuple(name for name, _parameter in named) | |
| if not self.parameter_names: | |
| raise ValueError("EMA teacher requires trainable student parameters") | |
| self._student = student | |
| self._parameters = tuple(parameter for _name, parameter in named) | |
| self._modules = tuple(student.modules()) | |
| self.shadow = { | |
| name: parameter.detach().float().clone() for name, parameter in named | |
| } | |
| self.native = { | |
| name: shadow.to(dtype=parameter.dtype).clone() | |
| for (name, parameter), shadow in zip( | |
| named, self.shadow.values(), strict=True | |
| ) | |
| } | |
| self._shadows = tuple(self.shadow.values()) | |
| self._native = tuple(self.native.values()) | |
| self.update_count = 0 | |
| self._active = False | |
| def _refresh_native(self) -> None: | |
| groups = defaultdict(lambda: ([], [])) | |
| for native, shadow in zip(self._native, self._shadows, strict=True): | |
| natives, shadows = groups[(native.device, native.dtype)] | |
| natives.append(native) | |
| shadows.append(shadow) | |
| for natives, shadows in groups.values(): | |
| for start in range(0, len(natives), 64): | |
| torch._foreach_copy_( | |
| natives[start : start + 64], shadows[start : start + 64] | |
| ) | |
| def _validate_student(self, student) -> None: | |
| if student is not self._student: | |
| raise RuntimeError("EMA teacher is bound to a different student model") | |
| if any(not parameter.requires_grad for parameter in self._parameters): | |
| raise RuntimeError("EMA student trainable parameter schema changed") | |
| def apply(self, student): | |
| """Temporarily evaluate ``student`` through EMA parameter storage.""" | |
| if self._active: | |
| raise RuntimeError("EMA teacher parameter context is not reentrant") | |
| self._validate_student(student) | |
| original_data = tuple(parameter.data for parameter in self._parameters) | |
| training_modes = tuple(module.training for module in self._modules) | |
| self._active = True | |
| try: | |
| student.eval() | |
| for parameter, native in zip( | |
| self._parameters, self._native, strict=True | |
| ): | |
| parameter.data = native | |
| with torch.no_grad(): | |
| yield student | |
| finally: | |
| for parameter, original in zip( | |
| self._parameters, original_data, strict=True | |
| ): | |
| parameter.data = original | |
| for module, training in zip(self._modules, training_modes, strict=True): | |
| module.training = training | |
| self._active = False | |
| def update(self, student) -> None: | |
| self._validate_student(student) | |
| for start in range(0, len(self._shadows), 64): | |
| shadows = self._shadows[start : start + 64] | |
| parameters = self._parameters[start : start + 64] | |
| torch._foreach_mul_(shadows, self.decay) | |
| torch._foreach_add_(shadows, parameters, alpha=1.0 - self.decay) | |
| self.update_count += 1 | |
| self._refresh_native() | |
| def state_dict(self) -> dict: | |
| return { | |
| "schema": self.STATE_SCHEMA, | |
| "decay": self.decay, | |
| "update_count": self.update_count, | |
| "parameter_names": list(self.parameter_names), | |
| "shadow": { | |
| name: value.detach().cpu() for name, value in self.shadow.items() | |
| }, | |
| } | |
| def is_finite(self) -> bool: | |
| norms = [] | |
| for start in range(0, len(self._shadows), 64): | |
| norms.extend( | |
| torch._foreach_norm(self._shadows[start : start + 64], float("inf")) | |
| ) | |
| return bool(torch.isfinite(torch.stack(norms)).all()) | |
| def load_state_dict(self, state: dict) -> None: | |
| if state.get("schema") not in { | |
| self.STATE_SCHEMA, | |
| *self.LEGACY_STATE_SCHEMAS, | |
| }: | |
| raise ValueError("unsupported EMA teacher state schema") | |
| if float(state["decay"]) != self.decay: | |
| raise ValueError( | |
| f"EMA decay mismatch: checkpoint={state['decay']}, current={self.decay}" | |
| ) | |
| if tuple(state["parameter_names"]) != self.parameter_names: | |
| raise ValueError("EMA parameter schema mismatch") | |
| if set(state["shadow"]) != set(self.parameter_names): | |
| raise ValueError("EMA shadow parameter keys mismatch") | |
| for name, value in state["shadow"].items(): | |
| if value.shape != self.shadow[name].shape: | |
| raise ValueError(f"EMA shadow shape mismatch for {name}") | |
| self.shadow[name].copy_(value.to(self.shadow[name])) | |
| self.update_count = int(state["update_count"]) | |
| self._refresh_native() | |
| __all__ = ["EMATeacher"] | |
Xet Storage Details
- Size:
- 5.7 kB
- Xet hash:
- ab2d996793d280f32620c515f5c570c58106431b707583dcc239845a8c9ceb38
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.