File size: 1,322 Bytes
784f572 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | """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()
|