rlhf_with_bridge.py

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
DISCLAIMER:
This example was generated by AI as a demonstration of RL adoption, showing how to
blend Hugging Face inference with Megatron-Bridge training. It is for educational
purposes only and not production-ready. Review, validate, and adapt before use.

Dummy RLHF example: Hugging Face inference + Megatron training via Bridge

Requirements
- CUDA-capable GPU(s)
- PyTorch with CUDA
- transformers >= 4.41
- Megatron-Core + Megatron-Bridge installed in this repo environment

What this shows
- HF policy model generates sampled continuations
- HF reward model (sentiment) scores responses
- Megatron policy (created via Bridge from the same HF model) is trained with a simple REINFORCE-style loss using the Megatron microbatch pipeline

Run (single GPU)
```bash
uv run python examples/rl/rlhf_with_bridge.py \
  --hf-policy-model Qwen/Qwen3-0.6B \
  --hf-reward-model distilbert-base-uncased-finetuned-sst-2-english \
  --train-iters 5 --mbs 1 --gbs 1 --seq-length 256 --max-new-tokens 32
```

Run (multi-GPU)
```bash
uv run python -m torch.distributed.run --nproc_per_node=2 examples/rl/rlhf_with_bridge.py \
  --hf-policy-model Qwen/Qwen3-0.6B \
  --hf-reward-model distilbert-base-uncased-finetuned-sst-2-english \
  --train-iters 20 --mbs 1 --gbs 2 --seq-length 256 --max-new-tokens 32
```

Notes
- This script is intentionally minimal and not a full PPO implementation.
- It requires CUDA; CPU-only runs are not supported by Megatron initialization.
- Replace the reward model or scoring logic as needed to test different reward functions.
"""

import argparse
import os
from dataclasses import dataclass
from typing import Iterable, Iterator

import torch
import torch.nn.functional as F
from megatron.core.pipeline_parallel import get_forward_backward_func
from megatron.core.process_groups_config import ProcessGroupCollection
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

from megatron.bridge import AutoBridge
from megatron.bridge.models.hf_pretrained.utils import is_safe_repo
from megatron.bridge.models.model_provider import get_model
from megatron.bridge.training.config import (
    CheckpointConfig,
    ConfigContainer,
    DistributedDataParallelConfig,
    LoggerConfig,
    OptimizerConfig,
    SchedulerConfig,
    TokenizerConfig,
    TrainingConfig,
)
from megatron.bridge.training.initialize import initialize_megatron, set_jit_fusion_options
from megatron.bridge.training.optim import setup_optimizer


@dataclass
class Args:
    """Dataclass configuration for training."""

    hf_policy_model: str
    hf_reward_model: str
    prompts: list[str]
    max_new_tokens: int
    learning_rate: float
    micro_batch_size: int
    global_batch_size: int
    train_iters: int
    seq_length: int
    trust_remote_code: bool = False


def build_config(provider, args: Args) -> ConfigContainer:
    """Build a config for the training."""
    provider.tensor_model_parallel_size = 1
    provider.pipeline_model_parallel_size = 1
    provider.context_parallel_size = 1
    provider.seq_length = args.seq_length
    provider.finalize()

    train = TrainingConfig(
        micro_batch_size=args.micro_batch_size,
        global_batch_size=args.global_batch_size,
        train_iters=args.train_iters,
    )

    optimizer = OptimizerConfig(
        optimizer="adam",
        lr=args.learning_rate,
        min_lr=args.learning_rate,
        weight_decay=0.0,
        adam_beta1=0.9,
        adam_beta2=0.95,
        use_distributed_optimizer=False,
        bf16=True,
    )

    scheduler = SchedulerConfig(
        lr_decay_style="constant",
        lr_warmup_iters=0,
        start_weight_decay=0.033,
        end_weight_decay=0.033,
        lr_decay_iters=args.train_iters,
        override_opt_param_scheduler=True,
    )

    ddp = DistributedDataParallelConfig()

    tokenizer = TokenizerConfig(
        tokenizer_type="HuggingFaceTokenizer",
        tokenizer_model=args.hf_policy_model,
    )

    checkpoint = CheckpointConfig(
        save_interval=0,
        save=None,
        load=None,
        async_save=False,
        fully_parallel_save=False,
        fully_parallel_load=False,
    )

    logger = LoggerConfig()

    cfg = ConfigContainer(
        model=provider,
        train=train,
        optimizer=optimizer,
        scheduler=scheduler,
        ddp=ddp,
        tokenizer=tokenizer,
        checkpoint=checkpoint,
        logger=logger,
        dataset=None,
    )
    cfg.validate()
    return cfg


def make_microbatch_iterator(batch: dict, num_microbatches: int) -> Iterator[dict]:
    """Make a microbatch iterator from a batch."""

    def _gen() -> Iterable[dict]:
        for _ in range(num_microbatches):
            yield batch

    return iter(_gen())


@torch.no_grad()
def refit_hf_from_megatron(
    bridge: AutoBridge,
    megatron_models: list,
    hf_model: AutoModelForCausalLM,
    *,
    show_progress: bool = False,
    cpu: bool = False,
) -> None:
    """Update the in-memory HF policy with current Megatron weights.

    This avoids writing to disk and lets the next rollout use the latest policy.
    """
    # Map parameter names to Tensors for fast lookup
    name_to_param = dict(hf_model.named_parameters())
    for name, tensor in bridge.export_hf_weights(megatron_models, cpu=cpu, show_progress=show_progress):
        param = name_to_param.get(name)
        if param is None:
            continue
        param.detach().copy_(tensor.to(param.device, dtype=param.dtype))


def build_parser() -> argparse.ArgumentParser:
    """Build a parser for the command line arguments."""
    p = argparse.ArgumentParser(description="Dummy RLHF: HF inference + Megatron training via Bridge")
    p.add_argument("--hf-policy-model", type=str, default="Qwen/Qwen3-0.6B")
    p.add_argument("--hf-reward-model", type=str, default="distilbert-base-uncased-finetuned-sst-2-english")
    p.add_argument("--trust-remote-code", action="store_true", help="if trust_remote_code")
    p.add_argument("--max-new-tokens", type=int, default=32)
    p.add_argument("--lr", type=float, default=5e-5)
    p.add_argument("--mbs", type=int, default=1)
    p.add_argument("--gbs", type=int, default=1)
    p.add_argument("--train-iters", type=int, default=5)
    p.add_argument("--seq-length", type=int, default=256)
    p.add_argument(
        "--prompts",
        type=str,
        nargs="*",
        default=[
            "Write a positive one-sentence movie review:",
            "Say something cheerful about the weather:",
        ],
    )
    return p


def main() -> None:
    """Main function."""
    parser = build_parser()
    ns = parser.parse_args()
    args = Args(
        hf_policy_model=ns.hf_policy_model,
        hf_reward_model=ns.hf_reward_model,
        prompts=list(ns.prompts),
        max_new_tokens=ns.max_new_tokens,
        learning_rate=ns.lr,
        micro_batch_size=ns.mbs,
        global_batch_size=ns.gbs,
        train_iters=ns.train_iters,
        seq_length=ns.seq_length,
        trust_remote_code=ns.trust_remote_code,
    )

    # Resolve per-rank device up front for multi-GPU runs
    if torch.cuda.is_available():
        local_rank = int(os.getenv("LOCAL_RANK", "0"))
        local_device = torch.device(f"cuda:{local_rank}")
    else:
        local_rank = -1
        local_device = torch.device("cpu")

    # HF tokenizer/model for generation (policy sampling) and HF reward pipeline
    hf_policy_model = args.hf_policy_model
    gen_tokenizer = AutoTokenizer.from_pretrained(
        hf_policy_model,
        trust_remote_code=is_safe_repo(
            trust_remote_code=args.trust_remote_code,
            hf_path=hf_policy_model,
        ),
    )
    # Use left padding for decoder-only models to avoid generation warnings and ensure correctness
    gen_tokenizer.padding_side = "left"
    if gen_tokenizer.pad_token_id is None:
        gen_tokenizer.pad_token = gen_tokenizer.eos_token
    hf_gen_model = AutoModelForCausalLM.from_pretrained(
        hf_policy_model,
        trust_remote_code=is_safe_repo(
            trust_remote_code=args.trust_remote_code,
            hf_path=hf_policy_model,
        ),
    )
    # Ensure pad_token_id is set on model config/generation config
    if getattr(hf_gen_model.config, "pad_token_id", None) is None:
        hf_gen_model.config.pad_token_id = gen_tokenizer.pad_token_id
    if (
        getattr(hf_gen_model, "generation_config", None) is not None
        and getattr(hf_gen_model.generation_config, "pad_token_id", None) is None
    ):
        hf_gen_model.generation_config.pad_token_id = gen_tokenizer.pad_token_id
    hf_gen_model.to(local_device)

    reward_pipe = pipeline(
        "sentiment-analysis",
        model=args.hf_reward_model,
        device=(local_rank if torch.cuda.is_available() else -1),
        truncation=True,
    )

    # Bridge: load HF, create Megatron provider and training stack
    bridge = AutoBridge.from_hf_pretrained(
        hf_policy_model,
        trust_remote_code=is_safe_repo(
            trust_remote_code=args.trust_remote_code,
            hf_path=hf_policy_model,
        ),
    )
    provider = bridge.to_megatron_provider(load_weights=True)

    cfg = build_config(provider, args)

    # Initialize Megatron (requires CUDA for real training environments)
    initialize_megatron(cfg=cfg)
    set_jit_fusion_options(cfg.model, cfg.train.micro_batch_size)

    # Get process group collection after initialization
    pg_collection = ProcessGroupCollection.use_mpu_process_groups()

    # Build model + optimizer + scheduler
    model_list = get_model(
        cfg.model,
        cfg.ddp,
        overlap_param_gather_with_optimizer_step=False,
        use_torch_fsdp2=cfg.dist.use_torch_fsdp2,
        data_parallel_random_init=cfg.rng.data_parallel_random_init,
        pg_collection=pg_collection,
    )
    model = model_list[0]
    optimizer, scheduler = setup_optimizer(
        optimizer_config=cfg.optimizer,
        scheduler_config=cfg.scheduler,
        model=model_list,
        use_gloo_process_groups=cfg.dist.use_gloo_process_groups,
    )

    forward_backward = get_forward_backward_func()

    # Simple loop: generate with HF, compute reward, train Megatron policy with REINFORCE
    device = local_device

    for step in range(args.train_iters):
        # 1) Generate responses with HF model
        batch_inputs = gen_tokenizer(
            args.prompts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=args.seq_length - args.max_new_tokens,
        ).to(device)

        with torch.no_grad():
            gen_out = hf_gen_model.generate(
                **batch_inputs,
                do_sample=True,
                top_p=0.9,
                temperature=0.7,
                max_new_tokens=args.max_new_tokens,
                pad_token_id=gen_tokenizer.pad_token_id,
            )

        # 2) Decode and score rewards with HF reward model
        texts = gen_tokenizer.batch_decode(gen_out, skip_special_tokens=True)
        rewards = []
        for t in texts:
            r = reward_pipe(t)[0]
            # reward: positive score if label == POSITIVE else 1 - score
            score = float(r["score"]) if r["label"].upper().startswith("POS") else float(1.0 - r["score"])
            rewards.append(score)
        rewards_t = torch.tensor(rewards, dtype=torch.float32, device=device)

        # 3) Build inputs for Megatron policy: we train on the generated full sequence
        # Teacher-forcing on generated sequence to get log-probs
        policy_inputs = gen_tokenizer(
            texts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=args.seq_length,
        )
        input_ids = policy_inputs["input_ids"].to(device)
        attention_mask = policy_inputs["attention_mask"].to(device)

        # Microbatch iterator yields one batch per microstep
        mb_iter = make_microbatch_iterator(
            {"input_ids": input_ids, "attention_mask": attention_mask},
            num_microbatches=1,
        )

        def rl_loss_fn(outputs: torch.Tensor, batch: dict) -> tuple[torch.Tensor, dict]:
            # outputs: [batch, seq, vocab] when TP=1
            logits = outputs
            # shift for next-token prediction
            shift_logits = logits[:, :-1, :]
            shift_labels = batch["input_ids"][:, 1:]
            # CE per position
            log_probs = F.log_softmax(shift_logits, dim=-1)
            # gather logp of taken tokens
            ll_selected = log_probs.gather(dim=-1, index=shift_labels.unsqueeze(-1)).squeeze(-1)
            # mask padding
            shift_mask = batch["attention_mask"][:, 1:].float()
            ll_selected = ll_selected * shift_mask
            # REINFORCE loss: -reward * sum_t logp(a_t)
            # broadcast rewards to sequence length
            r = rewards_t.view(-1, 1).expand_as(ll_selected)
            policy_loss = -(r * ll_selected).sum(dim=1).mean()
            return policy_loss, {"reinforce_loss": policy_loss.detach()}

        def forward_step_fn(data_iterator, megatron_model):
            batch = next(data_iterator)
            out = megatron_model(
                input_ids=batch["input_ids"],
                attention_mask=batch.get("attention_mask"),
                position_ids=batch.get("position_ids"),
            )
            return out, (lambda _o: rl_loss_fn(out, batch))

        model.train()
        forward_backward(
            forward_step_func=forward_step_fn,
            data_iterator=mb_iter,
            model=model,
            num_microbatches=1,
            seq_length=args.seq_length,
            micro_batch_size=args.micro_batch_size,
            decoder_seq_length=args.seq_length,
            forward_only=False,
        )
        optimizer.step()
        scheduler.step(increment=args.global_batch_size)

        # Refit updated Megatron policy weights back to the HF policy for next rollout
        refit_hf_from_megatron(bridge, model_list, hf_gen_model, show_progress=False, cpu=True)

        if (step + 1) % 1 == 0:
            print(f"Step {step + 1}/{args.train_iters} | mean reward: {rewards_t.mean().item():.4f}")


if __name__ == "__main__":
    # Helpful default when running locally without torchrun
    os.environ.setdefault("WORLD_SIZE", "1")
    main()