File size: 1,085 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 | """Adapter for the pinned CANDI hybrid discrete/continuous sampler."""
from dlb.adapters.base import BaseTeacherAdapter, CheckpointSelection
from dlb.runner import RunRequest
class CANDIAdapter(BaseTeacherAdapter):
identity = "dlb.adapters.candi:v1"
upstream = "candi"
supported_models = frozenset({"candi"})
teacher_families = {"candi": "hybrid_candi"}
batch_sizes = {("candi", "lm1b"): 2, ("candi", "owt"): 2}
def _loader_batch_size(self, request: RunRequest) -> int:
del request
return 16
def _sampling_overrides(
self, request: RunRequest, checkpoint: CheckpointSelection
) -> list[str]:
del request, checkpoint
return [
"algo=candi",
"algo.sampler=cached",
"algo.mixed_coeff=0.5",
"algo.step_size=1.0",
"algo.temp=1.0",
"algo.use_percentile_scheduling=True",
"sampling.noise_removal=ancestral",
"trainer.accumulate_grad_batches=1",
"eval.disable_ema=False",
]
adapter = CANDIAdapter()
|