| """Adapter for pinned FLM and FMLM checkpoints.""" | |
| from dlb.adapters.base import BaseTeacherAdapter, CheckpointSelection | |
| from dlb.runner import RunRequest | |
| class FLMAdapter(BaseTeacherAdapter): | |
| identity = "dlb.adapters.flm:v2" | |
| upstream = "flm" | |
| supported_models = frozenset({"flm", "fmlm"}) | |
| teacher_families = {"flm": "continuous_flm", "fmlm": "continuous_flm"} | |
| batch_sizes = { | |
| ("flm", "lm1b"): 32, | |
| ("flm", "owt"): 16, | |
| ("fmlm", "lm1b"): 16, | |
| ("fmlm", "owt"): 16, | |
| } | |
| def _sampling_overrides( | |
| self, request: RunRequest, checkpoint: CheckpointSelection | |
| ) -> list[str]: | |
| del checkpoint | |
| overrides = [ | |
| f"algo={request.model_id}", | |
| "algo.backbone=dit", | |
| "eval.disable_ema=False", | |
| "sampling.solver=euler", | |
| ] | |
| if request.model_id == "flm": | |
| overrides.extend(["algo.double_temb=False", "sampling.gamma=0.0"]) | |
| else: | |
| gamma = "0.8" if request.dataset_id == "lm1b" else "1.0" | |
| overrides.extend( | |
| [ | |
| "algo.double_temb=True", | |
| "algo.learnable_loss_weighting=False", | |
| f"sampling.gamma={gamma}", | |
| ] | |
| ) | |
| return overrides | |
| adapter = FLMAdapter() | |