{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "\n", "\n", "[![Documentation](https://img.shields.io/badge/Documentation-latest-blue.svg)](https://oumi.ai/docs/en/latest/index.html)\n", "[![Discord](https://img.shields.io/discord/1286348126797430814?label=Discord)](https://discord.gg/oumi)\n", "[![GitHub Repo stars](https://img.shields.io/github/stars/oumi-ai/oumi)](https://github.com/oumi-ai/oumi)\n", "
\n", "\n", "๐Ÿ‘‹ Welcome to Open Universal Machine Intelligence (Oumi)!\n", "\n", "๐Ÿš€ Oumi is a fully open-source platform that streamlines the entire lifecycle of foundation models - from [data preparation](https://oumi.ai/docs/en/latest/resources/datasets/datasets.html) and [training](https://oumi.ai/docs/en/latest/user_guides/train/train.html) to [evaluation](https://oumi.ai/docs/en/latest/user_guides/evaluate/evaluate.html) and [deployment](https://oumi.ai/docs/en/latest/user_guides/launch/launch.html). Whether you're developing on a laptop, launching large scale experiments on a cluster, or deploying models in production, Oumi provides the tools and workflows you need.\n", "\n", "๐Ÿค Make sure to join our [Discord community](https://discord.gg/oumi) to get help, share your experiences, and contribute to the project! If you are interested in joining one of the community's open-science efforts, check out our [open collaboration](https://oumi.ai/community) page.\n", "\n", "โญ If you like Oumi and you would like to support it, please give it a star on [GitHub](https://github.com/oumi-ai/oumi)." ] }, { "cell_type": "markdown", "metadata": { "id": "S7bYaH10SgtN" }, "source": [ "# OpenEnv GRPO with trl" ] }, { "cell_type": "markdown", "metadata": { "id": "QkhNGqE1SgtP" }, "source": [ "In this tutorial notebook, we're going to use Oumi to train an agentic model on an [OpenEnv](https://github.com/meta-pytorch/OpenEnv) Echo reinforcement learning (RL) environment with the GRPO algorithm. To achieve this, we use the trl library by Hugging Face with a custom rollout function to interact with the vLLM server and OpenEnv environment. This notebook is derived from trl's [Echo environment example](https://github.com/huggingface/trl/blob/main/examples/scripts/openenv/echo.py)." ] }, { "cell_type": "markdown", "metadata": { "id": "fHDr11SqSgtP" }, "source": [ "# ๐Ÿ“‹ Prerequisites\n", "\n", "โ—**NOTICE:** This notebook needs to be running on a machine with at least two GPUs.\n", "\n", "## Oumi Installation\n", "\n", "First, let's install the latest versions of Oumi and OpenEnv. You can find more detailed instructions [here](https://oumi.ai/docs/en/latest/get_started/installation.html)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!pip install uv && uv pip install \"oumi[gpu] @ git+https://github.com/oumi-ai/oumi.git\"\n", "!uv pip install git+https://github.com/meta-pytorch/OpenEnv.git" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "id": "JPmWKRVCSgtP" }, "outputs": [], "source": [ "import os\n", "from pathlib import Path\n", "\n", "tutorial_dir = \"openenv_tutorial\"\n", "\n", "Path(tutorial_dir).mkdir(parents=True, exist_ok=True)\n", "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\" # Disable warnings from HF." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Start OpenEnv and vLLM servers\n", "\n", "We need to run 2 servers in addition to the trl trainer. The OpenEnv server receives actions from the LLM and returns the updated state and reward. The vLLM server is used for inference, and updates it weights over training with the updated model weights from the trainer. We start these with separate subprocesses." ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Overwriting openenv_tutorial/start_openenv_server.py\n" ] } ], "source": [ "%%writefile $tutorial_dir/start_openenv_server.py\n", "\n", "import os\n", "import subprocess\n", "import sys\n", "import threading\n", "import time\n", "from pathlib import Path\n", "\n", "import requests\n", "\n", "\n", "def stream_output(pipe, prefix=\"\"):\n", " \"\"\"Stream output lines from subprocess pipe to stdout.\"\"\"\n", " for line in iter(pipe.readline, \"\"):\n", " print(f\"{prefix}{line}\", end=\"\")\n", " pipe.close()\n", "\n", "\n", "print(\"โšก Starting FastAPI server for Echo Environment...\")\n", "\n", "work_dir = str(Path.cwd().parent.absolute())\n", "\n", "server_process = subprocess.Popen(\n", " [\n", " sys.executable,\n", " \"-m\",\n", " \"uvicorn\",\n", " \"envs.echo_env.server.app:app\",\n", " \"--host\",\n", " \"0.0.0.0\",\n", " \"--port\",\n", " \"8001\",\n", " ],\n", " env={**os.environ, \"PYTHONPATH\": f\"{work_dir}/src\"},\n", " stdout=subprocess.PIPE,\n", " stderr=subprocess.PIPE,\n", " text=True,\n", " cwd=work_dir,\n", ")\n", "\n", "# Start background threads to stream errors\n", "threading.Thread(\n", " target=stream_output, args=(server_process.stderr, \"๐Ÿ”ฅ [stderr] \"), daemon=True\n", ").start()\n", "\n", "print(\"โณ Waiting for server to start...\")\n", "time.sleep(5)\n", "\n", "try:\n", " response = requests.get(\"http://0.0.0.0:8001/health\", timeout=2)\n", " print(\"\\nโœ… Echo Environment server is running!\")\n", "except Exception as e:\n", " print(f\"\\nโŒ Server failed to start: {e}\")\n", " print(\"\\n๐Ÿ“‹ Checking error output...\")\n", " server_process.poll()\n", " if server_process.stderr:\n", " stderr = server_process.stderr.read()\n", " if stderr:\n", " print(stderr)\n", " raise\n", "\n", "try:\n", " input(\"Press Enter to exit...\\n\")\n", "finally:\n", " print(\"๐Ÿ›‘ Stopping server...\")\n", " server_process.terminate()\n", " server_process.wait()" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Servers started. PIDs: 3787594 3787595\n" ] } ], "source": [ "import subprocess\n", "\n", "# Start both servers in the background\n", "server1 = subprocess.Popen(\n", " [\n", " \"bash\",\n", " \"-c\",\n", " (\n", " \"CUDA_VISIBLE_DEVICES=0 trl vllm-serve \"\n", " \"--model Qwen/Qwen2.5-0.5B-Instruct \"\n", " \"--log-level warning \"\n", " \"--host 0.0.0.0 --port 8000\"\n", " ),\n", " ]\n", ")\n", "server2 = subprocess.Popen([\"python\", f\"{tutorial_dir}/start_openenv_server.py\"])\n", "\n", "print(\"Servers started. PIDs:\", server1.pid, server2.pid)" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "โšก Starting FastAPI server for Echo Environment...\n", "โณ Waiting for server to start...\n", "โŒ Server not ready: HTTPConnectionPool(host='0.0.0.0', port=8000): Max retries exceeded with url: /health (Caused by NewConnectionError(': Failed to establish a new connection: [Errno 111] Connection refused'))\n", "๐Ÿ”ฅ [stderr] INFO: Started server process [3787596]\n", "๐Ÿ”ฅ [stderr] INFO: Waiting for application startup.\n", "๐Ÿ”ฅ [stderr] INFO: Application startup complete.\n", "๐Ÿ”ฅ [stderr] INFO: Uvicorn running on http://0.0.0.0:8001 (Press CTRL+C to quit)\n", "\n", "โœ… Echo Environment server is running!\n", "Press Enter to exit...\n", "โŒ Server not ready: HTTPConnectionPool(host='0.0.0.0', port=8000): Max retries exceeded with url: /health (Caused by NewConnectionError(': Failed to establish a new connection: [Errno 111] Connection refused'))\n", "โŒ Server not ready: HTTPConnectionPool(host='0.0.0.0', port=8000): Max retries exceeded with url: /health (Caused by NewConnectionError(': Failed to establish a new connection: [Errno 111] Connection refused'))\n", "INFO 10-31 17:31:32 [__init__.py:216] Automatically detected platform cuda.\n", "INFO 10-31 17:31:33 [utils.py:328] non-default args: {'disable_log_stats': True, 'worker_extension_cls': 'trl.scripts.vllm_serve.WeightSyncWorkerExtension', 'model_impl': 'vllm', 'model': 'Qwen/Qwen2.5-0.5B-Instruct'}\n", "โŒ Server not ready: HTTPConnectionPool(host='0.0.0.0', port=8000): Max retries exceeded with url: /health (Caused by NewConnectionError(': Failed to establish a new connection: [Errno 111] Connection refused'))\n", "INFO 10-31 17:31:41 [__init__.py:742] Resolved architecture: Qwen2ForCausalLM\n", "INFO 10-31 17:31:41 [__init__.py:1815] Using max model len 32768\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "`torch_dtype` is deprecated! Use `dtype` instead!\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "INFO 10-31 17:31:42 [scheduler.py:222] Chunked prefill is enabled with max_num_batched_tokens=16384.\n", "โŒ Server not ready: HTTPConnectionPool(host='0.0.0.0', port=8000): Max retries exceeded with url: /health (Caused by NewConnectionError(': Failed to establish a new connection: [Errno 111] Connection refused'))\n", "โŒ Server not ready: HTTPConnectionPool(host='0.0.0.0', port=8000): Max retries exceeded with url: /health (Caused by NewConnectionError(': Failed to establish a new connection: [Errno 111] Connection refused'))\n", "INFO 10-31 17:31:51 [__init__.py:216] Automatically detected platform cuda.\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:52 [core.py:654] Waiting for init message from front-end.\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:52 [core.py:76] Initializing a V1 LLM engine (v0.10.2) with config: model='Qwen/Qwen2.5-0.5B-Instruct', speculative_config=None, tokenizer='Qwen/Qwen2.5-0.5B-Instruct', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=32768, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=False, kv_cache_dtype=auto, device_config=cuda, decoding_config=DecodingConfig(backend='auto', disable_fallback=False, disable_any_whitespace=False, disable_additional_properties=False, reasoning_backend=''), observability_config=ObservabilityConfig(show_hidden_metrics_for_version=None, otlp_traces_endpoint=None, collect_detailed_traces=None), seed=0, served_model_name=Qwen/Qwen2.5-0.5B-Instruct, enable_prefix_caching=True, chunked_prefill_enabled=True, use_async_output_proc=True, pooler_config=None, compilation_config={\"level\":3,\"debug_dump_path\":\"\",\"cache_dir\":\"\",\"backend\":\"\",\"custom_ops\":[],\"splitting_ops\":[\"vllm.unified_attention\",\"vllm.unified_attention_with_output\",\"vllm.mamba_mixer2\",\"vllm.mamba_mixer\",\"vllm.short_conv\",\"vllm.linear_attention\",\"vllm.plamo2_mamba_mixer\",\"vllm.gdn_attention\"],\"use_inductor\":true,\"compile_sizes\":[],\"inductor_compile_config\":{\"enable_auto_functionalized_v2\":false},\"inductor_passes\":{},\"cudagraph_mode\":1,\"use_cudagraph\":true,\"cudagraph_num_of_warmups\":1,\"cudagraph_capture_sizes\":[512,504,496,488,480,472,464,456,448,440,432,424,416,408,400,392,384,376,368,360,352,344,336,328,320,312,304,296,288,280,272,264,256,248,240,232,224,216,208,200,192,184,176,168,160,152,144,136,128,120,112,104,96,88,80,72,64,56,48,40,32,24,16,8,4,2,1],\"cudagraph_copy_inputs\":false,\"full_cuda_graph\":false,\"pass_config\":{},\"max_capture_size\":512,\"local_cache_dir\":null}\n", "โŒ Server not ready: HTTPConnectionPool(host='0.0.0.0', port=8000): Max retries exceeded with url: /health (Caused by NewConnectionError(': Failed to establish a new connection: [Errno 111] Connection refused'))\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:52 [worker_base.py:595] Injected into for extended collective_rpc calls ['close_communicator', 'init_communicator', 'update_named_param']\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "[W1031 17:31:53.898867831 ProcessGroupNCCL.cpp:981] Warning: TORCH_NCCL_AVOID_RECORD_STREAMS is the default now, this environment variable is thus deprecated. (function operator())\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0\n", "[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0\n", "[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0\n", "[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0\n", "[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0\n", "[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:53 [parallel_state.py:1165] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, TP rank 0, EP rank 0\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m WARNING 10-31 17:31:53 [topk_topp_sampler.py:69] FlashInfer is not available. Falling back to the PyTorch-native implementation of top-p & top-k sampling. For the best performance, please install FlashInfer.\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:53 [gpu_model_runner.py:2338] Starting to load model Qwen/Qwen2.5-0.5B-Instruct...\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:53 [gpu_model_runner.py:2370] Loading model from scratch...\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:53 [cuda.py:362] Using Flash Attention backend on V1 engine.\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:53 [weight_utils.py:348] Using model weights format ['*.safetensors']\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:53 [weight_utils.py:406] No model.safetensors.index.json found in remote.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Loading safetensors checkpoint shards: 0% Completed | 0/1 [00:00: Failed to establish a new connection: [Errno 111] Connection refused'))\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:58 [backends.py:539] Using cache directory: /home/wizeng/.cache/vllm/torch_compile_cache/5d31f4c583/rank_0_0/backbone for vLLM's torch.compile\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:58 [backends.py:550] Dynamo bytecode transform time: 3.53 s\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:59 [backends.py:161] Directly load the compiled graph(s) for dynamic shape from the cache, took 1.361 s\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:31:59 [monitor.py:34] torch.compile takes 3.53 s in total\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:32:00 [gpu_worker.py:298] Available KV cache memory: 64.71 GiB\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:32:00 [kv_cache_utils.py:864] GPU KV cache size: 5,654,128 tokens\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:32:00 [kv_cache_utils.py:868] Maximum concurrency for 32,768 tokens per request: 172.55x\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Capturing CUDA graphs (mixed prefill-decode, PIECEWISE): 81%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ | 54/67 [00:01<00:00, 38.49it/s]" ] }, { "name": "stdout", "output_type": "stream", "text": [ "โŒ Server not ready: HTTPConnectionPool(host='0.0.0.0', port=8000): Max retries exceeded with url: /health (Caused by NewConnectionError(': Failed to establish a new connection: [Errno 111] Connection refused'))\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Capturing CUDA graphs (mixed prefill-decode, PIECEWISE): 100%|โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆ| 67/67 [00:01<00:00, 40.82it/s]\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:32:02 [gpu_model_runner.py:3118] Graph capturing finished in 2 secs, took 0.50 GiB\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:32:02 [gpu_worker.py:391] Free memory on device (78.59/79.19 GiB) on startup. Desired GPU memory utilization is (0.9, 71.27 GiB). Actual usage is 0.93 GiB for weight, 5.57 GiB for peak activation, 0.07 GiB for non-torch memory, and 0.5 GiB for CUDAGraph memory. Replace gpu_memory_utilization config with `--kv-cache-memory=68779733708` to fit into requested memory, or `--kv-cache-memory=76635612672` to fully utilize gpu memory. Current kv cache memory in use is 69478085324 bytes.\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:32:02 [core.py:218] init engine (profile, create kv cache, warmup model) took 8.68 seconds\n", "INFO 10-31 17:32:04 [llm.py:295] Supported_tasks: ['generate']\n", "INFO 10-31 17:32:04 [__init__.py:36] No IOProcessor plugins requested by the model\n", "โœ… vLLM server is healthy!\n" ] } ], "source": [ "import time\n", "\n", "import requests\n", "\n", "URL = \"http://0.0.0.0:8000/health\"\n", "\n", "\n", "def check_vllm_health():\n", " \"\"\"Checks if the vLLM server is healthy.\"\"\"\n", " try:\n", " response = requests.get(URL, timeout=3)\n", " if response.status_code == 200:\n", " print(\"โœ… vLLM server is healthy!\")\n", " return True\n", " else:\n", " print(f\"โš ๏ธ Server responded with {response.status_code}\")\n", " except requests.RequestException as e:\n", " print(f\"โŒ Server not ready: {e}\")\n", " return False\n", "\n", "\n", "max_retries = 24\n", "for attempt in range(1, max_retries + 1):\n", " if check_vllm_health():\n", " break\n", " time.sleep(5)\n", "else:\n", " print(f\"โŒ Failed to start vLLM server after {max_retries} attempts.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Train the model!\n", "\n", "By providing a custom rollout function to interact with the OpenEnv and vLLM servers, we can use trl to do agentic GRPO training. We also need to provide a reward function that processes the reward value output by the environment.\n", "\n", "The following script defines the custom rollout and reward functions and runs the trainer. We run it as a subprocess so that we can set `CUDA_VISIBLE_DEVICES` to not conflict with the vLLM server." ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Overwriting openenv_tutorial/train.py\n" ] } ], "source": [ "%%writefile $tutorial_dir/train.py\n", "\n", "import requests\n", "from envs.echo_env import EchoEnv\n", "from envs.echo_env.models import EchoAction\n", "\n", "from oumi.core.configs import TrainingConfig\n", "from oumi.core.registry import RegistryType, register\n", "from oumi.train import train\n", "\n", "\n", "@register(\"env_reward\", RegistryType.REWARD_FUNCTION)\n", "def reward_from_env(completions, **kwargs):\n", " \"\"\"Reward function that uses the environment reward.\"\"\"\n", " # Extract environment rewards from kwargs (propagated via extra_fields)\n", " env_rewards = kwargs.get(\"env_reward\", [])\n", " if env_rewards:\n", " return [float(reward) for reward in env_rewards]\n", " else:\n", " # Fallback if env_reward is not available\n", " return [0.0] * len(completions)\n", "\n", "\n", "@register(\"echo_env_vllm_rollout\", RegistryType.ROLLOUT_FUNCTION)\n", "def echo_env_vllm_rollout(\n", " prompts: list[str], args, processing_class\n", ") -> dict[str, list]:\n", " \"\"\"Custom rollout function that generates completions via vLLM server and computes environment rewards.\n", "\n", " Args:\n", " prompts: List of prompts to generate from\n", " args: GRPOConfig containing all sampling parameters\n", " processing_class: Tokenizer/processor for decoding completions\n", "\n", " Returns:\n", " Dict containing prompt_ids, completion_ids, logprobs, and env_reward\n", " \"\"\" # noqa: E501\n", " # 1. Generate completions via vLLM inference server (running on port 8000)\n", " payload = {\n", " \"prompts\": prompts,\n", " \"n\": args.num_generations,\n", " \"temperature\": args.temperature,\n", " \"top_p\": args.top_p,\n", " \"top_k\": -1 if args.top_k is None else args.top_k,\n", " \"min_p\": 0.0 if args.min_p is None else args.min_p,\n", " \"max_tokens\": args.max_completion_length,\n", " \"repetition_penalty\": args.repetition_penalty,\n", " }\n", " response = requests.post(\"http://0.0.0.0:8000/generate/\", json=payload)\n", "\n", " if response.status_code != 200:\n", " print(f\"Error response: {response.text}\")\n", "\n", " response.raise_for_status()\n", " result = response.json()\n", "\n", " completions_text = processing_class.batch_decode(\n", " result[\"completion_ids\"], skip_special_tokens=True\n", " )\n", "\n", " # 2. Step through the environment to get rewards\n", " client = EchoEnv(base_url=\"http://0.0.0.0:8001\")\n", " env_result = client.reset()\n", " env_rewards = []\n", " for msg in completions_text:\n", " env_result = client.step(EchoAction(message=msg))\n", " env_rewards.append(env_result.reward)\n", "\n", " # 3. Add environment rewards as extra field\n", " result[\"env_reward\"] = env_rewards\n", "\n", " return result\n", "\n", "\n", "config = TrainingConfig.from_yaml(\"openenv_tutorial/grpo_train.yaml\")\n", "train(config)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Finally, we define the YAML training config, and kick off training!\n", "\n", "To enable logging to Weights and Biases, uncomment the relevant line in the config below, and make sure to [set up wandb](https://oumi.ai/docs/en/latest/development/dev_setup.html#optional-set-up-weights-and-biases) on your machine." ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Overwriting openenv_tutorial/grpo_train.yaml\n" ] } ], "source": [ "%%writefile $tutorial_dir/grpo_train.yaml\n", "\n", "model:\n", " model_name: \"Qwen/Qwen2-0.5B-Instruct\"\n", " model_max_length: 2048\n", " torch_dtype_str: \"bfloat16\"\n", " attn_implementation: \"sdpa\"\n", "\n", "data:\n", " train:\n", " datasets:\n", " - dataset_name: \"trl-lib/ultrafeedback-prompt\"\n", " split: \"train\"\n", " sample_count: 100\n", "\n", "training:\n", " trainer_type: \"TRL_GRPO\"\n", " per_device_train_batch_size: 8\n", " gradient_accumulation_steps: 4\n", "\n", " reward_functions: [\"env_reward\"]\n", "\n", " ddp_find_unused_parameters: False\n", " optimizer: \"adamw_torch_fused\"\n", "\n", " grpo:\n", " use_vllm: True\n", " rollout_function: \"echo_env_vllm_rollout\"\n", " max_completion_length: 2048\n", "\n", " dataloader_num_workers: \"auto\"\n", " dataloader_prefetch_factor: 32\n", "\n", " num_train_epochs: 1\n", " logging_steps: 1\n", " log_model_summary: False\n", " output_dir: \"openenv_tutorial/echo_grpo\"\n", " # Uncomment to enable wandb logging\n", " # enable_wandb: True" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[2025-10-31 17:32:17,665][oumi][rank0][pid:3788270][MainThread][WARNING]][training_config.py:149] Ignored model.model_max_length=2048 parameter for trainer TrainerType.TRL_GRPO.\n", "[2025-10-31 17:32:17,666][oumi][rank0][pid:3788270][MainThread][INFO]][train.py:117] Creating training.output_dir: openenv_tutorial/echo_grpo...\n", "[2025-10-31 17:32:17,668][oumi][rank0][pid:3788270][MainThread][INFO]][train.py:119] Created training.output_dir absolute path: /home/wizeng/repos/oumi/notebooks/openenv_tutorial/echo_grpo\n", "[2025-10-31 17:32:17,669][oumi][rank0][pid:3788270][MainThread][INFO]][train.py:117] Creating training.telemetry_dir: openenv_tutorial/echo_grpo/telemetry...\n", "[2025-10-31 17:32:17,672][oumi][rank0][pid:3788270][MainThread][INFO]][train.py:119] Created training.telemetry_dir absolute path: /home/wizeng/repos/oumi/notebooks/openenv_tutorial/echo_grpo/telemetry\n", "[2025-10-31 17:32:17,675][oumi][rank0][pid:3788270][MainThread][INFO]][torch_utils.py:80] Torch version: 2.8.0+cu128. NumPy version: 1.26.4\n", "[2025-10-31 17:32:17,675][oumi][rank0][pid:3788270][MainThread][INFO]][torch_utils.py:88] CUDA version: 12.8 \n", "[2025-10-31 17:32:17,676][oumi][rank0][pid:3788270][MainThread][INFO]][torch_utils.py:91] CuDNN version: 90.8.0\n", "[2025-10-31 17:32:17,825][oumi][rank0][pid:3788270][MainThread][INFO]][torch_utils.py:124] CPU cores: 208 CUDA devices: 1\n", "device(0)='NVIDIA H100 80GB HBM3' Capability: (9, 0) Memory: [Total: 79.19GiB Free: 78.68GiB Allocated: 0.0GiB Cached: 0.0GiB]\n", "[2025-10-31 17:32:17,831][oumi][rank0][pid:3788270][MainThread][INFO]][train.py:154] Oumi version: 0.4.3.dev7+gee26267d5.d20251030\n", "[2025-10-31 17:32:17,838][oumi][rank0][pid:3788270][MainThread][INFO]][train.py:156] Git revision hash: 568abe14e6e819ef5e844360cb054bc34b5911eb\n", "[2025-10-31 17:32:17,881][oumi][rank0][pid:3788270][MainThread][INFO]][train.py:157] Git tag: None\n", "[2025-10-31 17:32:17,885][oumi][rank0][pid:3788270][MainThread][INFO]][train.py:165] Resolved 'training.dataloader_num_workers=auto' to 'training.dataloader_num_workers=2'\n", "[2025-10-31 17:32:17,911][oumi][rank0][pid:3788270][MainThread][INFO]][train.py:313] Training config saved to openenv_tutorial/echo_grpo/telemetry/training_config.yaml\n", "[2025-10-31 17:32:18,276][oumi][rank0][pid:3788270][MainThread][INFO]][models.py:544] Using the model's built-in chat template for model 'Qwen/Qwen2-0.5B-Instruct'.\n", "[2025-10-31 17:32:18,835][oumi][rank0][pid:3788270][MainThread][INFO]][models.py:260] Building model using device_map: auto (DeviceRankInfo(world_size=1, rank=0, local_world_size=1, local_rank=0))...\n", "[2025-10-31 17:32:18,861][oumi][rank0][pid:3788270][MainThread][INFO]][models.py:336] Using model class: to instantiate model.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "`torch_dtype` is deprecated! Use `dtype` instead!\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "[2025-10-31 17:32:19,171][oumi][rank0][pid:3788270][MainThread][INFO]][torch_utils.py:288] \n", "Model Parameters Summary:\n", "๐Ÿ”ข Total parameters: 494,032,768\n", "๐Ÿ”— Embedding parameters: 136,134,656\n", "๐ŸŽฏ Trainable parameters: 494,032,768\n", "๐Ÿ”’ Frozen parameters: 0 (0.00%)\n", "\n", "INFO 10-31 17:32:19 [__init__.py:216] Automatically detected platform cuda.\n", "[2025-10-31 17:32:19,942][oumi][rank0][pid:3788270][MainThread][INFO]][torch_profiler_utils.py:164] PROF: Torch Profiler disabled!\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/home/wizeng/repos/oumi/src/oumi/builders/training.py:70: UserWarning: You are importing from 'rollout_func', which is an experimental feature. This API may change or be removed at any time without prior notice. Silence this warning by setting environment variable TRL_EXPERIMENTAL_SILENCE=1.\n", " trainer = HuggingFaceTrainer(cls(*args, **kwargs, args=hf_args), processor)\n", "The model is already on multiple devices. Skipping the move to device specified in `args`.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Warning: The cache directory for DeepSpeed Triton autotune, /home/wizeng/.triton/autotune, appears to be on an NFS system. While this is generally acceptable, if you experience slowdowns or hanging when DeepSpeed exits, it is recommended to set the TRITON_CACHE_DIR environment variable to a non-NFS path.\n", "INFO 10-31 17:32:20 [__init__.py:1433] Found nccl from library libnccl.so.2\n", "INFO 10-31 17:32:20 [pynccl.py:70] vLLM is using nccl==2.27.3\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:32:20 [__init__.py:1433] Found nccl from library libnccl.so.2\n", "\u001b[1;36m(EngineCore_DP0 pid=3787957)\u001b[0;0m INFO 10-31 17:32:20 [pynccl.py:70] vLLM is using nccl==2.27.3\n", "[2025-10-31 17:32:21,330][oumi][rank0][pid:3788270][MainThread][INFO]][device_utils.py:343] GPU Metrics Before Training: GPU runtime info: NVidiaGpuRuntimeInfo(device_index=0, device_count=2, used_memory_mb=75593.0, temperature=33, fan_speed=None, fan_speeds=None, power_usage_watts=123.946, power_limit_watts=700.0, gpu_utilization=0, memory_utilization=0, performance_state=0, clock_speed_graphics=1980, clock_speed_sm=1980, clock_speed_memory=2619).\n", "[2025-10-31 17:32:21,330][oumi][rank0][pid:3788270][MainThread][INFO]][train.py:553] Training init time: 3.665s\n", "[2025-10-31 17:32:21,330][oumi][rank0][pid:3788270][MainThread][INFO]][train.py:554] Starting training... (TrainerType.TRL_GRPO, transformers: 4.57.1)\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.\n", " 0%| | 0/25 [00:00