docs: large tech edit

README.md

@@ -1,81 +1,94 @@
-# Nemo-Reinforcer: A Scalable and Efficient Post-Training Library for Models Ranging from tiny to >100B Parameters, scaling from 1 GPU to 100s
+# NVIDIA NeMo-Reinforcer: Scalable and Efficient Post-Training for Large Language Models
 
 <!-- markdown all in one -->
-- [Nemo-Reinforcer: A Scalable and Efficient Post-Training Library for Models Ranging from tiny to \>100B Parameters, scaling from 1 GPU to 100s](#nemo-reinforcer-a-scalable-and-efficient-post-training-library-for-models-ranging-from-tiny-to-100b-parameters-scaling-from-1-gpu-to-100s)
-  - [Features](#features)
-  - [Installation](#installation)
-  - [Quick start](#quick-start)
-    - [SFT](#sft)
-      - [Single Node](#single-node)
-      - [Multi-node](#multi-node)
-    - [GRPO](#grpo)
-      - [Single Node](#single-node-1)
-      - [Multi-node](#multi-node-1)
-  - [Cluster Start](#cluster-start)
-
-**Nemo-Reinforcer** is a scalable and efficient post-training library designed for models ranging from 1 GPU to thousands, and from tiny to over 100 billion parameters.
-
-What you can expect:
-
-- **Seamless integration with HuggingFace** for ease of use, allowing users to leverage a wide range of pre-trained models and tools.
-- **High-performance implementation with Megatron core**, supporting various parallelism techniques for large models (>100B) and large context lengths.
-- **Efficient resource management using Ray**, enabling scalable and flexible deployment across different hardware configurations.
-- **Flexibility** with a modular design that allows easy integration and customization.
-- **Comprehensive documentation** that is both detailed and user-friendly, with practical examples.
-
-## Features
-
-_✅ Available now | 🔜 Coming in v0.2_
-
-- ✅ **Fast Generation** - vLLM backend for optimized inference
-- ✅ **HuggingFace Integration** - Works with 1-8B models (Qwen1.5, Llama)
-- ✅ **Distributed Training** - FSDP support and Ray-based infrastructure
-- ✅ **Environment Support** - Support for multi-environment training.
-- ✅ **Learning Algorithms** - GRPO (Group Relative Policy Optimization) and SFT (Supervised Fine-Tuning)
-- ✅ **Worker Isolation** - Process isolation between RL Actors (no worries about global state)
-- 🔜 **Larger Model Support** - Native PyTorch support for models up to 70B parameters
-- 🔜 **Advanced Parallelism** - FSDP2, TP, SP, and sequence packing for efficient training
-- 🔜 **Environment Isolation** - Dependency isolation between components
-- 🔜 **DPO Algorithm** - Direct Preference Optimization for alignment
-
-## Installation
+
+**NVIDIA NeMo-Reinforcer** is a powerful and scalable post-training library designed to efficiently align and fine-tune large language models (LLMs), ranging from small models to those exceeding 100 billion parameters. It leverages distributed computing frameworks to scale seamlessly from a single GPU to hundreds, enabling efficient training and experimentation.
+
+This library provides a flexible and modular framework for implementing various post-training techniques, with a focus on Reinforcement Learning from Human Preferences (RLHF) and Supervised Fine-Tuning (SFT).
+
+## Table of Contents
+
+- [Key Features](#key-features)
+- [Install NeMo-Reinforcer](#installation)
+- [Quick Start](#quick-start)
+  - [Post-Train with Supervised Fine-Tuning (SFT)](#supervised-fine-tuning-sft)
+    - [Run Single Node](#single-node)
+    - [Run Multi-node](#multi-node)
+  - [Post-Train with Group Relative Policy Optimization (GRPO)](#group-relative-policy-optimization-grpo)
+    - [Run Single Node](#single-node-grpo)
+    - [Run Multi-node](#multi-node-grpo)
+- [Set Up Clusters](#cluster-setup)
+- [Documentation](#documentation)
+- [Contributing](#contributing)
+- [Licenses](#license)
+- [Citation](#citation)
+- [Support](#support)
+
+## Key Features
+
+_✅ Available Now | 🔜 Coming Soon (v0.2)_
+
+- ✅ **Fast Generation:** Utilizes vLLM backend for optimized inference during evaluation and rollout.
+- ✅ **HuggingFace Integration:** Seamlessly integrates with Hugging Face Transformers, supporting a wide range of pre-trained models (e.g., Qwen1.5, Llama models up to 8B parameters).
+- ✅ **Scalable Distributed Training:** Leverages Fully Sharded Data Parallelism (FSDP) and a Ray-based infrastructure for efficient multi-GPU and multi-node training.
+- ✅ **Multi-Environment Support:** Enables training across diverse environments and datasets.
+- ✅ **Reinforcement Learning Algorithms:** Implements Group Relative Policy Optimization (GRPO) for effective preference alignment.
+- ✅ **Supervised Fine-Tuning (SFT):** Supports standard supervised fine-tuning for instruction following and task adaptation.
+- ✅ **Worker Isolation:** Ensures process isolation between RL actors, preventing unintended global state interference.
+- 🔜 **Larger Model Support:** Native PyTorch support for models up to 70B parameters.
+- 🔜 **Advanced Parallelism Techniques:** Implementation of FSDP2, Tensor Parallelism (TP), Pipeline Parallelism (PP), and sequence packing for enhanced training efficiency.
+- 🔜 **Environment Isolation:** Provides dependency isolation between different components of the training pipeline.
+- 🔜 **Direct Preference Optimization (DPO):** Integration of the Direct Preference Optimization algorithm for more direct preference learning.
+
+## Install NeMo-Reinforcer
+
+For a streamlined setup, we recommend using `uv`. Ensure you have Python 3.12 or a compatible version installed.
 
 ```sh
-# For faster setup we use `uv`
+# Install uv for faster package management
 pip install uv
 
-# Specify a virtual env that uses Python 3.12
-uv venv -p python3.12.9 .venv
-# Install NeMo-Reinforcer with vllm
+# Create a virtual environment with Python 3.12
+uv venv -p python3.12 .venv
+
+# Activate the virtual environment (optional, but recommended for consistency)
+# source .venv/bin/activate  # On Linux/macOS
+# .venv\Scripts\activate  # On Windows
+
+# Install NeMo-Reinforcer with vLLM support
 uv pip install -e .[vllm]
-# Install NeMo-Reinforcer with dev/test dependencies
-uv pip install -e '.[dev,test]'
 
-# Use uv run to launch any runs. 
-# Note that it is recommended to not activate the venv and instead use `uv run` since
-# it ensures consistent environment usage across different shells and sessions.
+# To install with development and testing dependencies:
+# uv pip install -e '.[dev,test]'
+
+# Running scripts with `uv run` ensures a consistent environment.
 # Example: uv run python examples/run_grpo_math.py
 ```
 
-## Quick start
+**Important Notes:**
 
-**Reminder**: Don't forget to set your HF_HOME and WANDB_API_KEY (if needed). You'll need to do a `huggingface-cli login` as well for Llama models.
+- It is generally recommended **not to explicitly activate the virtual environment** when using `uv`. Instead, use `uv run <command>` to execute scripts within the managed environment. This helps maintain consistency across different shells and sessions.
+- Ensure you have the necessary CUDA drivers and PyTorch installed compatible with your hardware.
 
-### SFT
+## Quick Start
 
-We provide a sample SFT experiment that uses the [SQuAD dataset](https://rajpurkar.github.io/SQuAD-explorer/).
+Before running any experiments, remember to set your `HF_HOME` environment variable and your `WANDB_API_KEY` if you intend to use Weights & Biases for logging. For accessing Llama models, you might also need to log in using `huggingface-cli login`.
 
-#### Single Node
+## Post-Train with Supervised Fine-Tuning (SFT)
 
-The default SFT experiment is configured to run on a single GPU. To launch the experiment,
+We provide an example SFT experiment using the [SQuAD dataset](https://rajpurkar.github.io/SQuAD-explorer/).
+
+#### Run Single Node SFT
+
+The default SFT configuration is set to run on a single GPU. To start the experiment:
 
 ```sh
 uv run python examples/run_sft.py
 ```
 
-This trains `Llama3.2-1B` on one GPU using the SQUAD dataset.
+This command will fine-tune the `Llama3.2-1B` model on the SQuAD dataset using a single GPU.
 
-If you have access to more GPUs, you can update the experiment accordingly. To run on 8 GPUs, we update the cluster configuration. We also switch to an 8B Llama base model and increase the batch size:
+To utilize more GPUs on a single node, you can modify the cluster configuration and potentially adjust the model and batch size:
 
 ```sh
 uv run python examples/run_sft.py \
@@ -85,95 +98,122 @@ uv run python examples/run_sft.py \
   cluster.gpus_per_node=8
 ```
 
-Refer to [sft.yaml](examples/configs/sft.yaml) for a full list of parameters that can be overridden.
+For a comprehensive list of configurable parameters, refer to the [sft.yaml](https://www.google.com/search?q=examples/configs/sft.yaml) file.
 
-#### Multi-node
+#### Run Multi-node SFT
 
-For distributed training across multiple nodes:
+For distributed training across multiple compute nodes, ensure that the `UV_CACHE_DIR` is set to a shared directory accessible by all worker nodes before executing any `uv run` commands.
 
-Set `UV_CACHE_DIR` to a directory that can be read from all workers before running any uv run command.
 ```sh
 export UV_CACHE_DIR=/path/that/all/workers/can/access/uv_cache
 ```
 
+The following is an example Slurm script for launching a multi-node SFT experiment with the `Llama-3.1-8B` model:
+
 ```sh
-# Run from the root of NeMo-Reinforcer repo
-NUM_ACTOR_NODES=2
-# Add a timestamp to make each job name unique
+#!/bin/bash
+#SBATCH --nodes=2
+#SBATCH --account=YOUR_ACCOUNT
+#SBATCH --job-name=sft_llama8b_2nodes
+#SBATCH --partition=YOUR_PARTITION
+#SBATCH --time=4:00:00
+#SBATCH --gres=gpu:8
+
+NUM_ACTOR_NODES=$SLURM_JOB_NUM_NODES
 TIMESTAMP=$(date +%Y%m%d_%H%M%S)
 
-# SFT experiment uses Llama-3.1-8B model
-COMMAND="uv pip install -e .; uv run ./examples/run_sft.py --config examples/configs/sft.yaml cluster.num_nodes=2 cluster.gpus_per_node=8 checkpointing.checkpoint_dir='results/sft_llama8b_2nodes' logger.wandb_enabled=True logger.wandb.name='sft-llama8b'" \
-RAY_DEDUP_LOGS=0 \
-UV_CACHE_DIR=YOUR_UV_CACHE_DIR \
-CONTAINER=YOUR_CONTAINER \
-MOUNTS="$PWD:$PWD" \
-sbatch \
-    --nodes=${NUM_ACTOR_NODES} \
-    --account=YOUR_ACCOUNT \
-    --job-name=YOUR_JOBNAME \
-    --partition=YOUR_PARTITION \
-    --time=4:0:0 \
-    --gres=gpu:8 \
-    ray.sub
+COMMAND="uv pip install -e .; uv run ./examples/run_sft.py --config examples/configs/sft.yaml cluster.num_nodes=$NUM_ACTOR_NODES cluster.gpus_per_node=8 checkpointing.checkpoint_dir='results/sft_llama8b_2nodes' logger.wandb_enabled=True logger.wandb.name='sft-llama8b'"
+RAY_DEDUP_LOGS=0
+UV_CACHE_DIR=YOUR_UV_CACHE_DIR
+CONTAINER=YOUR_CONTAINER # Replace with your container if using one
+MOUNTS="$PWD:$PWD"
+
+srun --nodes=$NUM_ACTOR_NODES --ntasks-per-node=1 \
+  --gres=gpu:8 \
+  --job-name=${SLURM_JOB_NAME} \
+  bash -c "source .venv/bin/activate && ${COMMAND}" # If not using uv run directly
 ```
 
-### GRPO
+**Note:** Adapt the Slurm parameters (`--account`, `--partition`, `--job-name`, `--time`, `--gres`) according to your cluster configuration. You might need to adjust the command if you are not directly using `uv run` within the Slurm script.
+
+### Post-Train with Group Relative Policy Optimization (GRPO)
 
-We have a reference GRPO experiment config set up trained for math benchmarks using the [OpenInstructMath2](https://huggingface.co/datasets/nvidia/OpenMathInstruct-2) dataset.
+We provide a reference GRPO experiment configuration for training on math benchmarks using the [OpenInstructMath2](https://huggingface.co/datasets/nvidia/OpenMathInstruct-2) dataset.
 
-#### Single Node
+#### Run Single Node GRPO
 
-To run GRPO on a single GPU for `Llama-3.2-1B-Instruct`:
+To run the GRPO math example on a single GPU using the `Llama-3.2-1B-Instruct` model:
 
 ```sh
-# Run the GRPO math example using a 1B parameter model
 uv run python examples/run_grpo_math.py
 ```
 
-By default, this uses the configuration in `examples/configs/grpo_math_1B.yaml`. You can customize parameters with command-line overrides. For example, to run on 8 gpus,
+This command utilizes the default configuration specified in `examples/configs/grpo_math_1B.yaml`. You can override any parameter in the configuration file using command-line arguments. For instance, to run on 8 GPUs:
 
 ```sh
-# Run the GRPO math example using a 1B parameter model using 8 GPUs
 uv run python examples/run_grpo_math.py \
   cluster.gpus_per_node=8
 ```
 
-You can override any of the parameters listed in the yaml configuration file. For example,
+Here are more examples of overriding configuration parameters:
 
 ```sh
 uv run python examples/run_grpo_math.py \
   policy.model_name="Qwen/Qwen2-1.5B" \
   checkpointing.checkpoint_dir="results/qwen1_5b_math" \
   logger.wandb_enabled=True \
   logger.wandb.name="grpo-qwen1_5b_math" \
-  logger.num_val_samples_to_print=10 \
+  logger.num_val_samples_to_print=10
 ```
 
-#### Multi-node
+#### Run Multi-node GRPO
+
+The following is an example Slurm script for launching a multi-node GRPO experiment with the `Llama-3.1-8B-Instruct` model:
 
 ```sh
-# Run from the root of NeMo-Reinforcer repo
-NUM_ACTOR_NODES=2
-# Add a timestamp to make each job name unique
+#!/bin/bash
+#SBATCH --nodes=2
+#SBATCH --account=YOUR_ACCOUNT
+#SBATCH --job-name=grpo_llama8b_2nodes
+#SBATCH --partition=YOUR_PARTITION
+#SBATCH --time=4:00:00
+#SBATCH --gres=gpu:8
+
+NUM_ACTOR_NODES=$SLURM_JOB_NUM_NODES
 TIMESTAMP=$(date +%Y%m%d_%H%M%S)
 
-# grpo_math_8b uses Llama-3.1-8B-Instruct model
-COMMAND="uv pip install -e .; uv run ./examples/run_grpo_math.py --config examples/configs/grpo_math_8B.yaml cluster.num_nodes=2 checkpointing.checkpoint_dir='results/llama8b_2nodes' logger.wandb_enabled=True logger.wandb.name='grpo-llama8b_math'" \
-RAY_DEDUP_LOGS=0 \
-UV_CACHE_DIR=YOUR_UV_CACHE_DIR \
-CONTAINER=YOUR_CONTAINER \
-MOUNTS="$PWD:$PWD" \
-sbatch \
-    --nodes=${NUM_ACTOR_NODES} \
-    --account=YOUR_ACCOUNT \
-    --job-name=YOUR_JOBNAME \
-    --partition=YOUR_PARTITION \
-    --time=4:0:0 \
-    --gres=gpu:8 \
-    ray.sub
+COMMAND="uv pip install -e .; uv run ./examples/run_grpo_math.py --config examples/configs/grpo_math_8B.yaml cluster.num_nodes=$NUM_ACTOR_NODES checkpointing.checkpoint_dir='results/llama8b_2nodes' logger.wandb_enabled=True logger.wandb.name='grpo-llama8b_math'"
+RAY_DEDUP_LOGS=0
+UV_CACHE_DIR=YOUR_UV_CACHE_DIR
+CONTAINER=YOUR_CONTAINER # Replace with your container if using one
+MOUNTS="$PWD:$PWD"
+
+srun --nodes=$NUM_ACTOR_NODES --ntasks-per-node=1 \
+  --gres=gpu:8 \
+  --job-name=${SLURM_JOB_NAME} \
+  bash -c "source .venv/bin/activate && ${COMMAND}" # If not using uv run directly
 ```
 
-## Cluster Start
+**Note:** Adjust the Slurm directives according to your cluster setup.
+
+## Set Up Clusters
+
+For detailed instructions on how to set up and launch NeMo-Reinforcer on Slurm or Kubernetes clusters, please refer to the dedicated [Cluster Setup](https://www.google.com/search?q=docs/cluster.md) documentation.
+
+## Documentation
+
+Comprehensive documentation, including API references and more detailed explanations of concepts and functionalities, will be available soon. Stay tuned for updates\!
+
+## Contributing
+
+We welcome contributions to NeMo-Reinforcer\! Please see our [Contributing Guidelines](https://github.com/NVIDIA/NeMo/blob/stable/CONTRIBUTING.md) for more information on how to get involved.
+
+## Licenses
+
+NVIDIA NeMo-Reinforcer is licensed under the [Apache License 2.0](https://github.com/NVIDIA/NeMo?tab=Apache-2.0-1-ov-file#readme).
+
+NeMo is licensed under the [NVIDIA AI PRODUCT AGREEMENT](https://www.nvidia.com/en-us/agreements/enterprise-software/product-specific-terms-for-ai-products/). By pulling and using the container, you accept the terms and conditions of this license.
+
+## Support
 
-Please visit [Cluster Start](docs/cluster.md) for how to get started on Slurm or Kubernetes.
+For questions, bug reports, or feature requests, please open an issue on our [GitHub repository](https://github.com/NVIDIA/NeMo-Reinforcer/issues).

docs/design_docs/chat_datasets.md

@@ -1,8 +1,10 @@
 # Data Format
 
-## HuggingFace Chat Datasets
+This guide outlines the required data format for Hugging Face chat datasets and demonstrates how to use chat templates with Hugging Face tokenizers to add special tokens or task-specific information.
 
-HuggingFace chat datasets are expected to have the following structure: Each example in the dataset should be a dictionary with a `messages` key. `messages` should be a list of dictionaries, each with a `role` and `content` key. `role` is typically one of `system`, `user`, and `assistant`. For example:
+## Hugging Face Chat Datasets
+
+Hugging Face chat datasets are expected to have the following structure: Each example in the dataset should be a dictionary with a `messages` key. The `messages` should be a list of dictionaries, each with a `role` and `content` key. The `role` typically has one of the following values: `system`, `user`, and `assistant`. For example:
 
 ```json
 {
@@ -23,9 +25,9 @@ HuggingFace chat datasets are expected to have the following structure: Each exa
 }
 ```
 
-### Chat Templates
+## Chat Templates
 
-Formatting the data in this way allows us to take advantage of HuggingFace tokenizers' `apply_chat_template` functionality to combine the messages. Chat templates can be used to add special tokens or task-specific information to each example in the dataset. Refer to the [HuggingFace apply_chat_template documentation](https://huggingface.co/docs/transformers/main/en/chat_templating#applychattemplate) for details.
+Formatting the data with chat templates allows us to take advantage of the Hugging Face tokenizers' `apply_chat_template` functionality to combine the messages. Chat templates can be used to add special tokens or task-specific information to each example in the dataset. Refer to the [HuggingFace apply_chat_template documentation](https://huggingface.co/docs/transformers/main/en/chat_templating#applychattemplate) for details.
 
 By default, `apply_chat_template` attempts to apply the `chat_template` associated with the tokenizer. However, in some cases, users might want to specify their own chat template. Also, note that many tokenizers do not have associated `chat_template`s, in which case an explicit chat template is required. Users can specify an explicit chat template string using Jinja format and can pass that string to `apply_chat_template`. 
 The following is an example using a simple template which prepends a role header to each turn:
@@ -58,4 +60,4 @@ assert output == expected_output
 :hide:
 ```
 
-For more details on creating chat templates, refer to the [HuggingFace documentation](https://huggingface.co/docs/transformers/v4.34.0/en/chat_templating#how-do-i-create-a-chat-template).
+For more details on creating chat templates, refer to the [Hugging Face documentation](https://huggingface.co/docs/transformers/v4.34.0/en/chat_templating#how-do-i-create-a-chat-template).