File size: 5,262 Bytes
6b56858 | 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | """Run Di4C language distillation with project-controlled student initialization."""
from __future__ import annotations
from pathlib import Path
import sys
from typing import Any
_INIT_MODES = frozenset({"hf_small", "teacher", "scratch"})
_EVAL_GLOBALS = {
"__builtins__": {},
"__import__": __import__,
"abs": abs,
"float": float,
"int": int,
"len": len,
"max": max,
"min": min,
"round": round,
"str": str,
"sum": sum,
}
def _install_upstream_path() -> Path:
source = Path.cwd()
sdtt_src = source / "sdtt" / "src"
sdtt_package_root = sdtt_src / "sdtt"
if not (sdtt_package_root / "main.py").is_file():
raise RuntimeError(f"Di4C SDTT source is missing: {sdtt_src}")
for path in (source, sdtt_src, sdtt_package_root):
text = str(path)
if text not in sys.path:
sys.path.insert(0, text)
return source
def _student_init_mode(config: Any) -> str:
mode = config.get("dlb_student_init", "hf_small")
if mode not in _INIT_MODES:
raise ValueError(
f"unsupported dlb_student_init={mode!r}; expected one of {sorted(_INIT_MODES)}"
)
return str(mode)
def _eval_expression(expression: Any) -> Any:
return eval(str(expression), _EVAL_GLOBALS, {})
def _device_count() -> int:
try:
import torch
except ImportError:
return 0
return int(torch.cuda.device_count())
def _register_omegaconf_resolvers() -> None:
from omegaconf import OmegaConf
resolvers = {
"cwd": lambda: str(Path.cwd()),
"device_count": _device_count,
"div_up": lambda numerator, denominator: (
int(numerator) + int(denominator) - 1
)
// int(denominator),
"eval": _eval_expression,
}
for name, resolver in resolvers.items():
try:
OmegaConf.register_new_resolver(name, resolver, replace=True)
except TypeError:
OmegaConf.clear_resolver(name)
OmegaConf.register_new_resolver(name, resolver)
def _teacher_initialized_student(config: Any):
from sdtt.loading_utils import get_diffusion
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(config.tokenizer.name)
return get_diffusion(config, tokenizer)
def _hf_small_initialized_student(loss: str, round: int, config: Any):
if round not in range(1, 8):
raise ValueError(
f"Round value is too large: should be 1 <= round <= 7. Actual value: `{round}`"
)
if loss not in ("kld", "mse", "tvd"):
raise ValueError(f"Valid losses sizes: kld, mse, tvd. Actual value: `{loss}`")
from sdtt.core.distill.multi_round_sdtt import MultiRoundSDTT
revision = f"baselines_{loss}_step_{round * 10_000}"
# Upstream load_small_student uses student_as_teacher=True, which overwrites
# the configured teacher checkpoint with the HF small student backbone.
return MultiRoundSDTT.from_pretrained(
"jdeschena/sdtt",
revision,
config=config,
student_as_teacher=False,
)
def _scratch_initialized_student(config: Any):
from sdtt.models.loading_utils import get_backbone
model = _teacher_initialized_student(config)
model.backbone = get_backbone(config, model.vocab_size)
if hasattr(model.backbone, "is_di4c"):
model.backbone.is_di4c = True
model.init_ema()
return model
def _patch_student_loader() -> None:
import sdtt
original = sdtt.load_small_student
if getattr(original, "_dlb_student_init_patch", False):
return
def load_small_student(
loss: str = "kld",
round: int = 1,
config: Any | None = None,
**kwargs: Any,
):
if config is None:
try:
return original(loss=loss, round=round, config=config, **kwargs)
except TypeError:
if kwargs:
return original(loss=loss, round=round, config=config)
raise
mode = _student_init_mode(config)
if mode == "hf_small":
kwargs.pop("student_as_teacher", None)
if kwargs:
raise TypeError(
f"unexpected load_small_student keyword(s): {sorted(kwargs)}"
)
return _hf_small_initialized_student(loss=loss, round=round, config=config)
if mode == "teacher":
return _teacher_initialized_student(config)
if mode == "scratch":
return _scratch_initialized_student(config)
raise AssertionError(f"unreachable Di4C student init mode: {mode}")
load_small_student._dlb_student_init_patch = True # type: ignore[attr-defined]
sdtt.load_small_student = load_small_student
def main() -> int:
_install_upstream_path()
_register_omegaconf_resolvers()
import hydra
_patch_student_loader()
import sdtt.main as upstream_main
@hydra.main(
version_base=None,
config_path="../upstreams/di4c/sdtt/src/sdtt/configs",
config_name="config",
)
def run(config: Any) -> None:
upstream_main.main.__wrapped__(config)
return int(run() or 0)
if __name__ == "__main__":
raise SystemExit(main())
|