Coverage for transformer_lens/tools/model_registry/checkpoints.py: 94%

1"""Training-checkpoint label data for checkpointed model families. 

2 

3Canonical home for the checkpoint schedules previously defined in 

4``transformer_lens/loading_from_pretrained.py``. The schedules are frozen 

5historical artifacts of the published training runs. 

6""" 

7 

8import logging 

9 

10from .registry_io import resolve_model_alias 

11 

12# The steps for which there are checkpoints in the stanford crfm models 

13STANFORD_CRFM_CHECKPOINTS: list[int] = ( 

14 list(range(0, 100, 10)) 

15 + list(range(100, 2000, 50)) 

16 + list(range(2000, 20000, 100)) 

17 + list(range(20000, 400000 + 1, 1000)) 

18) 

19 

20# Linearly spaced checkpoints for Pythia models, taken every 1000 steps. 

21# Batch size 2,097,152 tokens, so checkpoints every 2.1B tokens 

22PYTHIA_CHECKPOINTS: list[int] = [0, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512] + list( 

23 range(1000, 143000 + 1, 1000) 

24) 

25# Pythia V1 has log-spaced early checkpoints (see line above), but V0 doesn't 

26PYTHIA_V0_CHECKPOINTS: list[int] = list(range(1000, 143000 + 1, 1000)) 

27 

28 

29def get_checkpoint_labels(model_name: str) -> tuple[list[int], str]: 

30 """Return (checkpoint labels, label type) for a checkpointed model family. 

31 

32 Covers the HF-revision-checkpointed families (Pythia, stanford-crfm). 

33 Raises ValueError for models without published checkpoint schedules. 

34 """ 

35 official_name = resolve_model_alias(model_name) or model_name 

36 if official_name.startswith("stanford-crfm/"): 

37 return STANFORD_CRFM_CHECKPOINTS, "step" 

38 if official_name.startswith("EleutherAI/pythia"): 

39 if "v0" in official_name: 

40 return PYTHIA_V0_CHECKPOINTS, "step" 

41 logging.warning( 

42 "Pythia models on HF were updated on 4/3/23! add '-v0' to model name to access the old models." 

43 ) 

44 return PYTHIA_CHECKPOINTS, "step" 

45 if official_name.startswith(("NeelNanda/", "ArthurConmy/", "Baidicoot/")): 

46 # Legacy TransformerLens repos store checkpoints as files, not 

47 # revisions: checkpoints/<name>_<step-or-tokens>.pth. 

48 import re 

49 

50 from huggingface_hub import HfApi 

51 

52 labels = sorted( 

53 int(m.group(1)) 

54 for f in HfApi().list_repo_files(official_name) 

55 if (m := re.match(r"checkpoints/.*_(\d+)\.pth", f)) 

56 ) 

57 if not labels: 57 ↛ 58line 57 didn't jump to line 58 because the condition on line 57 was never true

58 raise ValueError(f"Model {official_name} is not checkpointed.") 

59 return labels, ("token" if labels[-1] > 1e9 else "step") 

60 raise ValueError(f"Model {official_name} is not checkpointed.")