launch.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

import warnings
# Let's not worry about the future too much
warnings.simplefilter(action="ignore", category=FutureWarning)

import shutil
import threading
from logging import getLogger
import os
import socket
import time
import fairscale.nn.model_parallel.initialize as fs_init
from fairscale.nn.model_parallel import (
    get_model_parallel_group,
    get_model_parallel_rank,
)

from xlean.onlinerl.worker import entry
from xlean.fastgen.bulk.args import GenArgs, WorkerArgs
from xlean.fastgen import utils
from xlean.codegen.xlformers.src.data.tokenizer import build_tokenizer
from xlean.codegen.xlformers.src.params import ConfStore, cfg_from_cli
from xlean.codegen.xlformers.src.args import ModelArgs, ValidArgs
from xlean.codegen.xlformers.src.model.transformer import build_model
from xlean.onlinerl.onlinerl import OnlineRLContext
from xlean.onlinerl.args import ORLTrainerArgs, OnlineRLArgs
from xlean.onlinerl.train import start_trainer
from xlean.prover.args import EvalArgs as ProverEvalArgs
from xlean.codegen.xlformers.src.slurm import get_global_rank, get_world_size, init_torch_distributed
from xlean.utils import initialize_logger
from xlean.config import paths, get_tokenizer_path, get_data_path, get_checkpoint_path

logger = getLogger()


def initialize_excepthook():
    # shutdown upon thread exceptions
    def excepthook(args):
        raise args.exc_value

    threading.excepthook = excepthook


def initialize_onlinerl_distributed(args: ORLTrainerArgs, orl: OnlineRLContext) -> None:
    # initialize distributed mode / model parallel for online RL

    logger.info("Starting init of onlinerl...")
    num_workers = get_world_size() - args.orl.num_trainers
    orl.initialize(num_workers, args.model_parallel_size)
    logger.info("Done init of onlinerl.")

    logger.info("Starting init of torch.distributed...")
    args.slurm = init_torch_distributed()
    logger.info("Done init of torch.distributed.")

    logger.info("Starting init of model parallel...")
    fs_init.initialize_model_parallel(args.model_parallel_size)
    utils.mp.initialize(
        mp_size=args.model_parallel_size,
        mp_rank=get_model_parallel_rank(),
        mp_group=get_model_parallel_group(),
    )
    logger.info("Done init of model parallel.")


def sigchld_handler(signum, frame):
    logger.warning(f"Received SIGCHLD, ignoring")


def launch_onlinerl_worker(
        args: ORLTrainerArgs,
        orl: OnlineRLContext,\
     ) -> None:
    logger.info(f"Launching worker ...")
    tokenizer = build_tokenizer(args.tokenizer)
    tokenizer._stop_tokens = [
        tokenizer.eos_id,
        tokenizer.piece_to_id("<|effect|>"),
        tokenizer.piece_to_id("<|fim_suffix|>"),  # because of a past tokenizer.py bug
    ]  # stop decoding after the tactic

    # build model
    args.model.vocab_size = tokenizer.n_words
    model = build_model(
        args.model,
        dtype=args.dtype,
        fp32_reduce_scatter=args.fp32_reduce_scatter,
        reshard_after_forward=args.reshard_after_forward,
    )

    entry(orl, model, tokenizer, args)


def main(args: ORLTrainerArgs):
    # ctx = mp.get_context("fork")
    assert shutil.which("lake") is not None, f"command `lake` not found on {socket.gethostname()}"

    # logger.info(asdict(args))
    orl = OnlineRLContext(timeout=args.deadlock_timeout)
    initialize_onlinerl_distributed(args, orl)
    initialize_excepthook()

    inner_rank = get_global_rank()
    inner_size = get_world_size()

    print(f"{orl.rank=}: {orl.is_worker=} {inner_rank=} {inner_size=} PID {os.getpid()}")

    hostname = socket.gethostname()
    # rank=num_trainers is the first worker
    dispatch_host = orl.all_group.broadcast_object(obj=hostname, src=args.orl.num_trainers)
    print(f"BulkGen dispatch host: {dispatch_host}")
    args.orl.eval.bulkgen_host = dispatch_host
    args.orl.master_host = dispatch_host  # no need to keep apart atm

    initialize_logger()

    if orl.is_worker:
        launch_onlinerl_worker(args, orl)
    else:
        logger.info(f"Launching trainers ...")
        start_trainer(args, orl)


if __name__ == "__main__":
    from xlean.codegen.xlformers.src.args import TokenizerArgs
    from xlean.codegen.xlformers.src.data.tokenizer import Tokenizer

    assert "GITHUB_ACCESS_TOKEN" in os.environ, "Set GITHUB_ACCESS_TOKEN environment variable"

    if "cl_toplang_128k_lean" not in ConfStore:
        ConfStore["cl_toplang_128k_lean"] = TokenizerArgs(
            tokenizer_cls="TiktokenTokenizer",
            model="cl_toplang_128k.tiktoken",
            num_reserved_special_tokens=256,
            basic_special_tokens=[
                "<|begin_of_text|>",
                "<|end_of_text|>",
                "<|fim_prefix|>",
                "<|fim_middle|>",
                "<|fim_suffix|>",
            ],
            extra_special_tokens=["<|effect|>"],
            pat_str=Tokenizer.DEFAULT_TIKTOKEN_PATTERN,
        )

    job_name = f"debughtps_{time.strftime('%m%d-%H%M', time.localtime())}"
    ConfStore["default_orl"] = OnlineRLArgs(
        num_trainers=8,
        eval=ProverEvalArgs(
            data=get_data_path("minif2f", "val.jsonl"),
            tokenizer=get_tokenizer_path("cl_toplang_128k.tiktoken"),
            use_multisample=True,
            max_expansions=20,
            timeout=0,
            tactic_timeout=120,
        ),
        bulk=WorkerArgs(
            gen=GenArgs(
                max_batch=32,
                max_seq=4096,
                max_gen=200,
                logprobs=True,
            )
        ),
        provable_id=16,  # "1" in Llama3 tokenizer
        unprovable_id=15,  # "0" in Llama3 tokenizer
    )

    args = ORLTrainerArgs(
        dump_dir="debug_run",
        orl=ConfStore["default_orl"],
        dump_freq=50,
        eval_freq=50,
        log_freq=1,
        valid=ValidArgs(
            ppl_files_str=get_data_path("leandojo_benchmark_4_v10", "random", "sft", "val.jsonl"),
            batch_size=4,
            seq_len=4096,
        ),
        finetuning_dir=get_checkpoint_path("lean4_v10", "1e-4_acc"),
        tokenizer=ConfStore["cl_toplang_128k_lean"],
        model=ModelArgs(
            dim=4096,
            n_layers=32,
            n_heads=32,
            n_kv_heads=8,
            ffn_dim_multiplier=1.3,
            rope_theta=500000.0,
            multiple_of=1024,
            efficient_attn="auto",
            vocab_parallel=True
        ),
        batch_size=4,
        seq_len=4096,
        deadlock_timeout=2 * 3600,  # 2h
        disable_wandb=True,
        iter_type="orl",
        enable_loss_tracker=True,
        gpu_check_level=-1,  # gpu check after each eval (-1 to disable)
        periodic_gpu_check=False,
    )

    args: ORLTrainerArgs = cfg_from_cli(base_config=args)
    main(args)