| """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}" |
| |
| |
| 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 |
| 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()) |
|
|
