update

Author: Lu-Yang666

examples/experimental/dapo/greso_dapo.py

@@ -0,0 +1,348 @@
+import os
+import sys
+from copy import deepcopy
+
+import torch.distributed as dist
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+from areal.api.alloc_mode import AllocationMode
+from areal.api.cli_args import GRPOConfig, load_expr_config
+from areal.api.io_struct import FinetuneSpec, StepInfo, WeightUpdateMeta
+from areal.dataset import get_custom_dataset
+from areal.engine.ppo.actor import FSDPPPOActor
+from areal.engine.sglang_remote import RemoteSGLangEngine
+from areal.platforms import current_platform
+from areal.utils import seeding, stats_tracker
+from areal.utils.data import (
+    broadcast_tensor_container,
+    cycle_dataloader,
+    tensor_container_to,
+)
+from areal.utils.device import log_gpu_stats
+from areal.utils.evaluator import Evaluator
+from areal.utils.hf_utils import load_hf_tokenizer
+from areal.utils.recover import RecoverHandler
+from areal.utils.saver import Saver
+from areal.utils.stats_logger import StatsLogger
+from areal.workflow.rlvr import RLVRWorkflow
+
+from typing import TYPE_CHECKING, Optional
+from datasets import load_dataset
+from datasets.distributed import split_dataset_by_node
+if TYPE_CHECKING:
+    from datasets import Dataset
+    from transformers.processing_utils import ProcessorMixin
+    from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
+
+def gsm8k_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **kwargs):
+    from areal.reward.math_parser import process_results
+
+    return int(process_results(completions, answer)[0])
+
+def load_greso_dataset(
+    path: str,
+    rank: int,
+    world_size: int,
+    type: str = "sft",
+    split: Optional[str] = None,
+    max_length: Optional[int] = None,
+    tokenizer: Optional["PreTrainedTokenizerFast"] = None,
+    processor: Optional["ProcessorMixin"] = None,
+    **kwargs,
+) -> "Dataset":
+    dataset = load_dataset("parquet", data_dir=path, split=split)
+
+    def process(sample):
+        return {"messages": sample["messages"], "answer": sample["answer"]}
+    
+    dataset = dataset.map(process)
+
+    # Filter out sequences longer than max_length if tokenizer and max_length are provided
+    if max_length is not None:
+
+        def filter_length(sample):
+            # Tokenize the user content to check length
+            content = sample["messages"][0]["content"]
+            tokens = tokenizer.encode(content)
+            return len(tokens) <= max_length
+
+        dataset = dataset.filter(filter_length)
+
+    dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size)
+    return dataset
+
+
+def main(args):
+    config, _ = load_expr_config(args, GRPOConfig)
+    config: GRPOConfig
+
+    rank = int(os.getenv("RANK"))
+    tokenizer = load_hf_tokenizer(config.tokenizer_path)
+
+    seeding.set_random_seed(config.seed, key=f"trainer{rank}")
+    allocation_mode = AllocationMode.from_str(config.allocation_mode)
+    parallel_strategy = allocation_mode.train
+    assert parallel_strategy is not None
+
+    # Initialize train engine
+    actor = FSDPPPOActor(config=config.actor)
+    actor.create_process_group(parallel_strategy=parallel_strategy)
+
+    train_dataset = load_greso_dataset(
+        path=config.train_dataset.path,
+        rank=actor.data_parallel_rank,
+        world_size=actor.data_parallel_world_size,
+        split="train",
+        max_length=config.train_dataset.max_length,
+        type=config.train_dataset.type,
+        tokenizer=tokenizer,
+    )
+    valid_dataset = load_greso_dataset(
+        path=config.valid_dataset.path,
+        rank=actor.data_parallel_rank,
+        world_size=actor.data_parallel_world_size,
+        split="test",
+        max_length=config.valid_dataset.max_length,
+        type=config.valid_dataset.type,
+        tokenizer=tokenizer,
+    )
+
+    # Create dataset and dataloaders
+    train_dataloader = StatefulDataLoader(
+        train_dataset,
+        batch_size=config.train_dataset.batch_size // actor.data_parallel_world_size,
+        shuffle=config.train_dataset.shuffle,
+        num_workers=config.train_dataset.num_workers,
+        collate_fn=lambda x: x,
+        drop_last=config.train_dataset.drop_last,
+    )
+    valid_dataloader = StatefulDataLoader(
+        valid_dataset,
+        batch_size=config.valid_dataset.batch_size // actor.data_parallel_world_size,
+        shuffle=config.valid_dataset.shuffle,
+        num_workers=config.valid_dataset.num_workers,
+        collate_fn=lambda x: x,
+        drop_last=config.valid_dataset.drop_last,
+    )
+    ft_spec = FinetuneSpec(
+        total_train_epochs=config.total_train_epochs,
+        dataset_size=len(train_dataloader) * config.train_dataset.batch_size,
+        train_batch_size=config.train_dataset.batch_size,
+    )
+
+    # Initialize inference engine
+    rollout = RemoteSGLangEngine(config.rollout)
+    rollout.initialize(train_data_parallel_size=parallel_strategy.dp_size)
+    eval_rollout = RemoteSGLangEngine(deepcopy(config.rollout))
+    # NOTE: eval does not have any offpolicyness control
+    eval_rollout.config.max_head_offpolicyness = int(1e12)
+    eval_rollout.initialize()
+
+    actor.initialize(None, ft_spec)
+    ref = None
+    if config.actor.kl_ctl > 0 and config.ref is not None:
+        ref = FSDPPPOActor(config=config.ref)
+        ref.create_process_group(parallel_strategy=parallel_strategy)
+        ref.initialize(None, ft_spec)
+
+    # NOTE: Weight update meta only requires address and free port of rank 0,
+    # but `WeightUpdateMeta.from_fsdp_xccl` has to be executed on all ranks
+    # due to `engine.get_param_specs()`.
+    # Therefore, we create weight update meta on all ranks, then broadcast the one on rank 0.
+    weight_update_meta = [
+        WeightUpdateMeta.from_fsdp_xccl(
+            AllocationMode.from_str(config.allocation_mode), actor
+        )
+    ]
+    dist.broadcast_object_list(weight_update_meta, src=0)
+    weight_update_meta = weight_update_meta[0]
+
+    # Create rollout workflow
+    if tokenizer.pad_token_id not in config.gconfig.stop_token_ids:
+        config.gconfig.stop_token_ids.append(tokenizer.pad_token_id)
+    if tokenizer.eos_token_id not in config.gconfig.stop_token_ids:
+        config.gconfig.stop_token_ids.append(tokenizer.eos_token_id)
+    workflow = RLVRWorkflow(
+        reward_fn=gsm8k_reward_fn,
+        gconfig=config.gconfig,
+        tokenizer=tokenizer,
+        enable_thinking=False,
+        dump_dir=os.path.join(
+            StatsLogger.get_log_path(config.stats_logger), "generated"
+        ),
+    )
+    eval_workflow = RLVRWorkflow(
+        reward_fn=gsm8k_reward_fn,
+        gconfig=config.gconfig.new(temperature=0.6),
+        tokenizer=tokenizer,
+        enable_thinking=False,
+        rollout_stat_scope="eval-rollout",
+        dump_dir=os.path.join(
+            StatsLogger.get_log_path(config.stats_logger), "generated-eval"
+        ),
+    )
+
+    # Run training.
+    saver = Saver(config.saver, ft_spec)
+    stats_logger = StatsLogger(config.stats_logger, ft_spec)
+    evaluator = Evaluator(config.evaluator, ft_spec)
+
+    recover_handler = RecoverHandler(config.recover, ft_spec)
+    recover_info = recover_handler.load(
+        actor,
+        saver,
+        evaluator,
+        stats_logger,
+        train_dataloader,
+        inference_engine=rollout,
+        weight_update_meta=weight_update_meta,
+    )
+    start_step = (
+        recover_info.last_step_info.next().global_step
+        if recover_info is not None
+        else 0
+    )
+
+    total_epochs = config.total_train_epochs
+    steps_per_epoch = len(train_dataloader)
+    max_steps = total_epochs * steps_per_epoch
+
+    data_generator = cycle_dataloader(train_dataloader)
+    for global_step in range(start_step, max_steps):
+        epoch = global_step // steps_per_epoch
+        step = global_step % steps_per_epoch
+        step_info = StepInfo(
+            global_step=global_step,
+            epoch=epoch,
+            epoch_step=step,
+            steps_per_epoch=steps_per_epoch,
+        )
+
+        with stats_tracker.record_timing("rollout"):
+            batch = None
+            if actor.is_data_parallel_head():
+                if config.async_training:
+                    batch = rollout.prepare_batch(
+                        train_dataloader,
+                        workflow=workflow,
+                        should_accept=lambda sample: True,
+                    )
+                else:
+                    batch = rollout.rollout_batch(
+                        next(data_generator),
+                        workflow=workflow,
+                        should_accept=lambda sample: True,
+                    )
+                batch = tensor_container_to(batch, actor.device)
+            batch = broadcast_tensor_container(
+                batch,
+                src_rank=actor.current_data_parallel_head(),
+                group=actor.context_and_model_parallel_group,
+            )
+        # Create barrier to synchronize all rollout processes.
+        dist.barrier(device_ids=[actor.device.index])
+        current_platform.synchronize()
+
+        if config.actor.recompute_logprob or config.actor.use_decoupled_loss:
+            with stats_tracker.record_timing("recompute_logp"):
+                logp = actor.compute_logp(batch)
+                batch["prox_logp"] = logp
+                log_gpu_stats("recompute logp")
+
+        if ref is not None:
+            with stats_tracker.record_timing("ref_logp"):
+                batch["ref_logp"] = ref.compute_logp(batch)
+                log_gpu_stats("ref logp")
+
+        with stats_tracker.record_timing("compute_advantage"):
+            actor.compute_advantages(batch)
+            log_gpu_stats("compute advantages")
+
+        with (
+            stats_tracker.record_timing("train_step"),
+            stats_tracker.scope("grpo_actor"),
+        ):
+            stats = actor.ppo_update(batch)
+            actor.step_lr_scheduler()
+            log_gpu_stats("ppo update")
+
+        # pause inference for updating weights, save, and evaluation
+        rollout.pause()
+
+        with stats_tracker.record_timing("update_weights"):
+            if dist.get_rank() == 0:
+                future = rollout.update_weights(weight_update_meta)
+            actor.upload_weights(weight_update_meta)
+            if dist.get_rank() == 0:
+                future.result()
+            dist.barrier(device_ids=[actor.device.index])
+            current_platform.synchronize()
+
+            actor.set_version(global_step + 1)
+            rollout.set_version(global_step + 1)
+            eval_rollout.set_version(global_step + 1)
+
+        with stats_tracker.record_timing("save"):
+            saver.save(actor, epoch, step, global_step, tokenizer=tokenizer)
+
+        with stats_tracker.record_timing("checkpoint_for_recover"):
+            recover_handler.dump(
+                actor,
+                step_info,
+                saver,
+                evaluator,
+                stats_logger,
+                train_dataloader,
+                tokenizer=tokenizer,
+            )
+
+        dist.barrier(device_ids=[actor.device.index])
+        current_platform.synchronize()
+
+        with stats_tracker.record_timing("eval"):
+
+            def evaluate_fn():
+                if actor.is_data_parallel_head():
+                    # Stats are logged in workflow
+                    # and will be exported later
+                    cnt = 0
+                    for data in valid_dataloader:
+                        for item in data:
+                            eval_rollout.submit(item, eval_workflow)
+                            cnt += 1
+                    eval_rollout.wait(cnt, timeout=None)
+                dist.barrier(device_ids=[actor.device.index])
+                current_platform.synchronize()
+
+            evaluator.evaluate(
+                evaluate_fn,
+                epoch,
+                step,
+                global_step,
+            )
+
+        dist.barrier(device_ids=[actor.device.index])
+        current_platform.synchronize()
+
+        # Upload statistics to the logger (e.g., wandb)
+        stats[0].update(
+            stats_tracker.export_all(reduce_group=actor.data_parallel_group)
+        )
+        stats_logger.commit(epoch, step, global_step, stats)
+
+        dist.barrier(device_ids=[actor.device.index])
+        current_platform.synchronize()
+
+        # Resume rollout
+        rollout.resume()
+
+    stats_logger.close()
+    eval_rollout.destroy()
+    rollout.destroy()
+    if ref is not None:
+        ref.destroy()
+    actor.destroy()
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

scripts/dapo.sh

@@ -1,16 +1,18 @@
 #!/usr/bin/env bash
 set -euo pipefail
-export CUDA_VISIBLE_DEVICES=0,1,2
+export CUDA_VISIBLE_DEVICES=1,2
 N_GPU=2
-EXP_NAME=gsm8k-dapo
-TRIAL_NAME=trial1
+EXP_NAME=greso-dapo
+TRIAL_NAME=trial0
 FILE_ROOT=/data/yl/AReaL/tmp/areal/experiments
 ACTOR_PATH=/data/yl/model/Qwen/Qwen2.5-1.5B-Instruct
+TRAIN_DATASET_PATH=/data/yl/dataset/greso
+VALID_DATASET_PATH=/data/yl/dataset/greso
 
 TOTAL_TRAIN_EPOCHS=1
 
 python3 -m areal.launcher.local \
-    examples/experimental/dapo/gsm8k_dapo.py \
+    examples/experimental/dapo/greso_dapo.py \
     --config examples/experimental/dapo/gsm8k_dapo.yaml \
     experiment_name="$EXP_NAME" \
     trial_name="$TRIAL_NAME" \
@@ -19,7 +21,12 @@ python3 -m areal.launcher.local \
     cluster.n_nodes=1 \
     cluster.n_gpus_per_node="$N_GPU" \
     cluster.fileroot="$FILE_ROOT" \
+    +gconfig.top_p=0.7 \
     actor.path="$ACTOR_PATH" \
     actor.optimizer.lr=1e-6 \
-    actor.optimizer.weight_decay=0.1 \
-    actor.overlong_reward_penalty=false 
+    actor.optimizer.weight_decay=0.01 \
+    actor.overlong_reward_penalty=false \
+    actor.ppo_n_minibatches=64 \
+    +actor.c_clip=10.0 \
+    train_dataset.path="$TRAIN_DATASET_PATH" \
+    valid_dataset.path="$VALID_DATASET_PATH"