Text Generation
Transformers
Safetensors
OpenPeerLLM
English
Arabic
riemann_openpeer
amd
ntk
code
fsdp
rocm
coding
mi325x
decoder
code gen
directml
openpeer
ryzen-ai
reasoning
rewriting
multimodal
paraphrase
prerelease
ntk-trainer
growth-ladder
style-transfer
text-rewriting
riemann-computing
progressive-growth
custom_code
Eval Results (legacy)
File size: 2,440 Bytes
db3982a | 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 | from __future__ import annotations
import argparse
import logging
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from transformers import AutoTokenizer
from riemann.growth_ladder import get_stage, next_stage
from riemann.growth import load_checkpoint_and_grow
from riemann.paths import CHECKPOINT_ROOT
logger = logging.getLogger("riemann.grow_and_continue")
TOKENIZER_NAME = "gpt2"
def main() -> None:
logging.basicConfig(
level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
)
parser = argparse.ArgumentParser(
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
)
parser.add_argument(
"--from-stage", required=True, help="the stage you just finished training"
)
parser.add_argument(
"--to-stage",
default=None,
help="target stage name (default: the next one in the ladder)",
)
parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT))
args = parser.parse_args()
from_stage = get_stage(args.from_stage)
to_stage = (
get_stage(args.to_stage) if args.to_stage else next_stage(from_stage.name)
)
from_dir = Path(args.checkpoint_root) / from_stage.name
to_dir = Path(args.checkpoint_root) / to_stage.name
if not (from_dir / "config.json").exists():
raise FileNotFoundError(
f"No checkpoint found at {from_dir} -- train this stage first"
)
logger.info("Growing %s -> %s", from_stage.name, to_stage.name)
grown_model = load_checkpoint_and_grow(str(from_dir), to_stage)
to_dir.mkdir(parents=True, exist_ok=True)
grown_model.save_pretrained(str(to_dir))
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
tokenizer.save_pretrained(str(to_dir))
n_params = sum((p.numel() for p in grown_model.parameters()))
logger.info(
"Saved grown checkpoint to %s (%d params, ~%.2fB)",
to_dir,
n_params,
n_params / 1000000000.0,
)
if not to_stage.feasible_on_this_machine:
logger.warning(
"%s is marked not feasible on this machine; expect this to be slow or run out of memory",
to_stage.name,
)
logger.info("Next: python scripts/pretrain.py --stage %s --resume", to_stage.name)
if __name__ == "__main__":
main()
|