File size: 1,505 Bytes
11c11f8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
945c5bf
11c11f8
 
 
 
 
 
 
 
 
 
 
 
 
 
945c5bf
11c11f8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from .benchmark import benchmark_hyper, run_baseline, run_hyper
from .common import (
    DEFAULT_SCALE_PRESETS,
    apply_standard_config_tweaks,
    build_model_from_args,
    cosine_lr,
    save_final_checkpoint,
    save_training_checkpoint,
    setup_cpu_runtime,
)
from .datasets import (
    GrowLengthDataset,
    PreTokenizedDataset,
    SequenceTokenDataset,
    build_sequence_dataset,
    build_token_buffer,
    format_dataset_example,
    matches_category_filter,
)
from .hyper import (
    GrowLengthScheduler,
    ProgressiveLoopScheduler,
    ProgressiveUnfreezer,
    SeedReplayMeZO,
    apply_reservoir_freezing,
    patch_training_loops,
)
from .loops import train_fast_loop, train_hyper_loop, train_standard_loop
from .optimizers import MeZOOptimizer

__all__ = [
    "DEFAULT_SCALE_PRESETS",
    "GrowLengthDataset",
    "GrowLengthScheduler",
    "MeZOOptimizer",
    "PreTokenizedDataset",
    "ProgressiveLoopScheduler",
    "ProgressiveUnfreezer",
    "SeedReplayMeZO",
    "SequenceTokenDataset",
    "benchmark_hyper",
    "build_sequence_dataset",
    "build_token_buffer",
    "format_dataset_example",
    "matches_category_filter",
    "apply_reservoir_freezing",
    "apply_standard_config_tweaks",
    "build_model_from_args",
    "cosine_lr",
    "patch_training_loops",
    "save_final_checkpoint",
    "save_training_checkpoint",
    "setup_cpu_runtime",
    "run_baseline",
    "run_hyper",
    "train_fast_loop",
    "train_hyper_loop",
    "train_standard_loop",
]