start script

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.<config_name>` and for inference via `--inference.<config_name>`.
+
+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
 

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)
 

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)

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()))

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()))

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,
         }
 
 

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()