diff --git a/README.md b/README.md index f0faac0c8d..2c97d2cb37 100644 --- a/README.md +++ b/README.md @@ -72,6 +72,19 @@ and then do the training on it uv run torchrun --nproc_per_node=2 src/zeroband/train.py @ configs/training/150M/A40.toml --data.path data/fake_rollout ``` +## RL launcher + +rl launcher is a script that allow to start training and inference at the same time and assign GPUs to each process. + +Under the hood its just start script a bit like torchrun do. + +```bash +uv run src/zeroband/rl_launcher.py --n_gpus 2 --train @ configs/training/debug.toml --train.optim.total_steps 10000 --inference @ configs/inference/debug.toml --inference.max_samples 10000 +``` + +You can pass any config that you would pass for training via `--train.` and for inference via `--inference.`. + +In the future this launcher will make sure that both training and inference configs are compatible with each other. For now there is no specific config validation logic. ## Checkpoints management diff --git a/src/zeroband/inference.py b/src/zeroband/inference.py index edeb5db98d..fcce7b700d 100644 --- a/src/zeroband/inference.py +++ b/src/zeroband/inference.py @@ -5,6 +5,7 @@ from pydantic_config import BaseConfig, parse_argv import vllm +from zeroband.logger import get_logger from zeroband.models import ModelName, name_to_hf_model from datasets import load_dataset @@ -19,6 +20,7 @@ class Config(BaseConfig): sample_per_file: int = 1024 max_samples: int | None = None output_path: str = "outputs" + tp: int = 1 @model_validator(mode="after") def validate_bs_and_sample_per_file(self): @@ -102,7 +104,8 @@ def get_parquet_table(generated_tokens: list[vllm.RequestOutput], step: int) -> def main(config: Config): # -> list[dict[str, Any]]: prompts = ["Write me a novel" for _ in range(5)] - llm = LLM(model=name_to_hf_model[config.name_model]) + llm = LLM(model=name_to_hf_model[config.name_model], tensor_parallel_size=config.tp) + logger = get_logger("INFERENCE") # tokenizer = llm.get_tokenizer() sampling_params = SamplingParams(temperature=0.7, top_p=0.95, max_tokens=100, presence_penalty=0.1, frequency_penalty=0.1) @@ -125,7 +128,9 @@ def main(config: Config): # -> list[dict[str, Any]]: # Get tokenized inputs prompts = fake_chat_template(messages) - generated_tokens = llm.generate(prompts, sampling_params) + generated_tokens = llm.generate(prompts, sampling_params, use_tqdm=False) + + logger.info(f"Generated {len(prompts)} prompts") table = get_parquet_table(generated_tokens, step) diff --git a/src/zeroband/logger.py b/src/zeroband/logger.py index 5c3e20c471..4e70e52a98 100644 --- a/src/zeroband/logger.py +++ b/src/zeroband/logger.py @@ -11,19 +11,19 @@ def __init__(self, local_rank: int): self.local_rank = local_rank def format(self, record): - log_format = "{asctime} [{levelname}] [Rank {local_rank}] {message}" + log_format = "{asctime} [{levelname}] [{name}] [Rank {local_rank}] {message}" formatter = logging.Formatter(log_format, style="{", datefmt="%H:%M:%S") - record.local_rank = self.local_rank # Add this line to set the local rank in the record + record.local_rank = self.local_rank return formatter.format(record) -def get_logger(config=None, name: str | None = None) -> logging.Logger: - global logger # Add this line to modify the global logger variable +def get_logger(name: str = "TRAIN") -> logging.Logger: + global logger if logger is not None: return logger world_info = get_world_info() - logger = logging.getLogger(name or __name__) + logger = logging.getLogger(name) if world_info.local_rank == 0: logger.setLevel(level=logging.INFO) diff --git a/src/zeroband/rl_launcher.py b/src/zeroband/rl_launcher.py new file mode 100644 index 0000000000..9d765b6f91 --- /dev/null +++ b/src/zeroband/rl_launcher.py @@ -0,0 +1,168 @@ +import atexit +import os +import signal +import sys + +from pydantic_config import parse_argv, BaseConfig +import torch +from zeroband.train import Config as TrainConfig +from zeroband.inference import Config as InferenceConfig +from zeroband.inference import main as inference +from zeroband.train import train +from zeroband.logger import get_logger + +import torch.multiprocessing as mp + +processes = [] + + +class Config(BaseConfig): + train: TrainConfig + inference: InferenceConfig + + n_gpus: int | None = None + ratio: float = 0.5 # for now we have half train and half inference + + torchrun_rdzv_address: str = "localhost" + torchrun_rdzv_port: int = 29500 + + +class EnvWrapper: + """ + This class wrapp a function call and overide the environment variables + FYI: cannot use a simple function because of pickle issues + """ + + def __init__(self, fn, envs): + self.fn = fn + self.envs = envs + + def __call__(self, *args, **kwargs): + os.environ.update(self.envs) + return self.fn(*args, **kwargs) + + +def _cuda_available_devices(gpus_ids: list[int]) -> str: + return ",".join(map(str, gpus_ids)) + + +def train_torchrun(config: TrainConfig, rdzv_address: str, rdzv_port: int, gpus_ids: list[int]) -> list[mp.Process]: + """ + This funciton simulated torchrun but manage to wrap a function call instead of starting from a files. + + Under the hood it just created n_proc processes and set the environment variables for each of them. + + Torchrun is doing this as well under the hood but wrap logs and more advance rdzv features that we don't need yet. + + """ + # Set start method to 'spawn' to avoid CUDA initialization issues + config.gpus_ids = gpus_ids + nproc_per_node = len(gpus_ids) + + processes = [] + for rank in range(nproc_per_node): + # Prepare environment variables + envs = {} + envs["MASTER_ADDR"] = rdzv_address + envs["MASTER_PORT"] = str(rdzv_port) + envs["CUDA_VISIBLE_DEVICES"] = _cuda_available_devices(gpus_ids) + envs["RANK"] = str(rank) + envs["LOCAL_RANK"] = str(rank) + envs["LOCAL_WORLD_SIZE"] = str(nproc_per_node) + envs["WORLD_SIZE"] = str(nproc_per_node) + fn_env = EnvWrapper(train, envs) + p = mp.Process(target=fn_env, args=(config,)) + p.start() + processes.append(p) + + return processes + + +def inference_run(config: InferenceConfig, gpus_ids: list[int]) -> list[mp.Process]: + """ + This function is used to run inference by creating a sub process. + """ + envs = {"CUDA_VISIBLE_DEVICES": _cuda_available_devices(gpus_ids)} + + config.tp = len(gpus_ids) + + fn_env = EnvWrapper(inference, envs) + process = mp.Process(target=fn_env, args=(config,)) + process.start() + + return [process] + + +def cleanup_subprocesses(): + """Kill all registered multiprocessing processes""" + for process in processes: + try: + if process.is_alive(): # Check if mp.Process is still running + logger.info(f"Terminating process with PID {process.pid}") + process.terminate() # Try to terminate gracefully + + # Wait for a bit to see if it terminates + process.join(timeout=3) + + # If it's still alive, force kill it + if process.is_alive(): + logger.info(f"Process {process.pid} didn't terminate, killing...") + # On Unix, we can use os.kill for a hard kill + try: + os.kill(process.pid, signal.SIGKILL) + except Exception as e: + logger.info(f"Failed to kill process {process.pid}: {e}") + except Exception as e: + logger.info(f"Error cleaning up process: {e}") + + +def signal_handler(sig, frame): + logger.info(f"Received signal {sig}, cleaning up...") + cleanup_subprocesses() + sys.exit(0) + + +def main(config: Config): + if config.n_gpus is None: + config.n_gpus = torch.cuda.device_count() + + gpus_ids = list(range(config.n_gpus)) + cutoff = int(config.n_gpus * config.ratio) + + train_gpus_ids = gpus_ids[cutoff:] + inference_gpus_ids = gpus_ids[:cutoff] + + logger.info(f"start rl training with {len(train_gpus_ids)} GPUs, {len(inference_gpus_ids)}. Total: {len(gpus_ids)}") + logger.info(f"train_gpus_ids: {train_gpus_ids}") + logger.info(f"inference_gpus_ids: {inference_gpus_ids}") + + mp.set_start_method("spawn", force=True) + + train_processes = train_torchrun( + config.train, rdzv_address=config.torchrun_rdzv_address, rdzv_port=config.torchrun_rdzv_port, gpus_ids=train_gpus_ids + ) + inference_process = inference_run(config.inference, gpus_ids=inference_gpus_ids) + + processes.extend(train_processes) + processes.extend(inference_process) + + try: + for p in processes: + p.join() + except Exception as e: + logger.info(f"Error in main process: {e}") + # The cleanup will happen via atexit + + +if __name__ == "__main__": + # Register cleanup function to be called when the program exits + atexit.register(cleanup_subprocesses) + + # Register signal handlers to handle SIGINT and SIGTERM + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + + #### the code above alow to kill all subprocess like torchrun do + + logger = get_logger("RL_LAUNCHER") + main(Config(**parse_argv())) diff --git a/src/zeroband/train.py b/src/zeroband/train.py index f2b8e89b6f..32b80b4ba8 100644 --- a/src/zeroband/train.py +++ b/src/zeroband/train.py @@ -8,6 +8,7 @@ from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy # type: ignore import wandb + from zeroband.models import ModelName, get_model_and_tokenizer from zeroband.training.checkpoint import TrainingProgress, load_checkpoint_fsdp_state, save_checkpoint_fsdp_state from zeroband.training.data import DataConfig, get_dataloader @@ -20,7 +21,7 @@ from pydantic_config import BaseConfig, parse_argv from jaxtyping import Float, Int -from zeroband.training.world_info import get_world_info +from zeroband.training.world_info import WorldInfo, get_world_info class AdamConfig(BaseConfig): @@ -65,6 +66,8 @@ class Config(BaseConfig): optim: OptimConfig = OptimConfig() train: TrainConfig + gpus_ids: list[int] | None = None + @model_validator(mode="after") def check_batch_size(self): if self.data.batch_size is None: @@ -75,7 +78,7 @@ def check_batch_size(self): return self -def get_gradient_accumulation_steps(batch_size: int, micro_bs: int, data_workers: int) -> int: +def get_gradient_accumulation_steps(batch_size: int, micro_bs: int, data_workers: int, world_info: WorldInfo) -> int: assert batch_size % world_info.local_world_size == 0 batch_size = batch_size // world_info.local_world_size @@ -100,10 +103,38 @@ def apply_fsdp(model: torch.nn.Module, reshard_after_forward: bool): fully_shard(model, mp_policy=mp_policy, reshard_after_forward=reshard_after_forward) +def get_device_placement(gpus_ids: list[int] | None, world_info: WorldInfo) -> int: + """handle using a subset of GPUs. Should work like the CUDA_VISIBLE_DEVICES env var. + The reason we use this is because in the rl launcher, torch is initialized before the env var is set, so we cannot use the CUDA_VISIBLE_DEVICES env var. + """ + if gpus_ids is None: + return world_info.local_rank + + if world_info.local_rank >= len(gpus_ids): + raise ValueError(f"Local rank {world_info.local_rank} is greater than the number of available GPUs ({len(gpus_ids)})") + + return gpus_ids[world_info.local_rank] + + def train(config: Config): + logger = get_logger() + world_info = get_world_info() + + logger.info(f"start training on {world_info.world_size}") + + # Allow eager fallback during production so that that the training runs dont die + # However, in development, we want to know that we broke torch compile + torch._dynamo.config.suppress_errors = "ZERO_BAND_DEV" not in os.environ # type: ignore + torch.set_float32_matmul_precision("high") + torch.manual_seed(42) + + torch.cuda.set_device(get_device_placement(config.gpus_ids, world_info)) + # batch_size is the total batch size for all GPUs - gradient_accumulation_steps = get_gradient_accumulation_steps(config.optim.batch_size, config.train.micro_bs, config.data.num_workers) + gradient_accumulation_steps = get_gradient_accumulation_steps( + config.optim.batch_size, config.train.micro_bs, config.data.num_workers, world_info + ) model, tokenizer = get_model_and_tokenizer(config.name_model) @@ -200,16 +231,4 @@ def train(config: Config): if __name__ == "__main__": - # Allow eager fallback during production so that that the training runs dont die - # However, in development, we want to know that we broke torch compile - torch._dynamo.config.suppress_errors = "ZERO_BAND_DEV" not in os.environ # type: ignore - torch.set_float32_matmul_precision("high") - torch.manual_seed(42) - - config = Config(**parse_argv()) # type: ignore - world_info = get_world_info() - logger = get_logger() - - torch.cuda.set_device(world_info.local_rank) - - train(config) + train(Config(**parse_argv())) diff --git a/src/zeroband/training/world_info.py b/src/zeroband/training/world_info.py index a232306ce7..f132e872c1 100644 --- a/src/zeroband/training/world_info.py +++ b/src/zeroband/training/world_info.py @@ -18,18 +18,8 @@ def __init__(self): self.local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", 1)) self.nnodes = self.world_size // self.local_world_size - self.global_unique_id = os.environ.get("GLOBAL_UNIQUE_ID", None) - self.global_addr = os.environ.get("GLOBAL_ADDR", None) - self.global_port = int(os.environ.get("GLOBAL_PORT")) if "GLOBAL_PORT" in os.environ else None - self.global_world_size = int(os.environ.get("GLOBAL_WORLD_SIZE", 1)) - self.global_rank = int(os.environ.get("GLOBAL_RANK", 0)) - def __repr__(self): - return f"WorldInfo(world_size={self.world_size}, rank={self.rank}, local_rank={self.local_rank}, local_world_size={self.local_world_size}, nnodes={self.nnodes}, global_unique_id={self.global_unique_id}, global_addr={self.global_addr}, global_port={self.global_port}, global_world_size={self.global_world_size}, global_rank={self.global_rank})" - - @property - def diloco_rank(self): - return self.global_rank + return f"WorldInfo(world_size={self.world_size}, rank={self.rank}, local_rank={self.local_rank}, local_world_size={self.local_world_size}, nnodes={self.nnodes}, device_placement={self.device_placement})" def json(self) -> dict[str, int | str]: return { @@ -38,11 +28,6 @@ def json(self) -> dict[str, int | str]: "local_rank": self.local_rank, "local_world_size": self.local_world_size, "nnodes": self.nnodes, - "global_unique_id": self.global_unique_id, - "global_addr": self.global_addr, - "global_port": self.global_port, - "global_world_size": self.global_world_size, - "global_rank": self.global_rank, } diff --git a/tests/test_inference.py b/tests/test_inference.py index d06f2c7268..8822f3c0f8 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -14,11 +14,12 @@ def _test_torchrun(config, extra_args=[]): process = subprocess.Popen(cmd) result = process.wait() if result != 0: - pytest.fail(f"Process {result} failed {result}") + pytest.fail(f"Process failed {result}") -def test_inference(tmp_path): - _test_torchrun(config="inference/debug.toml", extra_args=["--output_path", str(tmp_path)]) +@pytest.mark.parametrize("tp", [1, 2]) +def test_inference(tmp_path, tp): + _test_torchrun(config="inference/debug.toml", extra_args=["--output_path", str(tmp_path), "--tp", str(tp)]) assert tmp_path.joinpath("step_0").exists() diff --git a/tests/test_train.py b/tests/test_train.py index f777bc52a9..cd662ce630 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -14,7 +14,7 @@ def _test_torchrun(num_gpus, config, extra_args=[]): process = subprocess.Popen(cmd) result = process.wait() if result != 0: - pytest.fail(f"Process {result} failed {result}") + pytest.fail(f"Process failed {result}") @pytest.mark.parametrize("num_gpus", [1, 2]) @@ -28,3 +28,14 @@ def test_train_with_rollout_file(fake_rollout_files_dir): """ path = fake_rollout_files_dir(steps=list(range(5)), num_files=8, batch_size=8, seq_len=16) # there is more file than batch here _test_torchrun(num_gpus=1, config="debug.toml", extra_args=["--data.path", str(path), "--no-data.fake"]) + + +def test_rl_launcher(): + """ + this test will create a fake rollout file and then train with it + """ + cmd = "uv run src/zeroband/rl_launcher.py --n_gpus 2 --train @ configs/training/debug.toml --inference @ configs/inference/debug.toml".split() + process = subprocess.Popen(cmd) + result = process.wait() + if result != 0: + pytest.fail(f"Process failed {result}")