Skip to content

[Tunix] Add agentic mode support to grpo_demo_llama3_qwen2.py. #920

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[Tunix] Add agentic mode support to grpo_demo_llama3_qwen2.py. #920

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 43 additions & 9 deletions scripts/grpo_demo_llama3_qwen2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,6 @@
from tunix.models.qwen2 import model as qwen2_lib
from tunix.models.qwen2 import params as qwen2_params
from tunix.rl import rl_cluster as rl_cluster_lib
from tunix.rl.grpo import grpo_learner
from tunix.rl.rollout import base_rollout
from tunix.sft import metrics_logger
from tunix.sft import profiler
Expand Down Expand Up @@ -243,11 +242,36 @@
required=False,
help="Name of dataset, required when data_source is tfds",
)

parser.add_argument(
"--agentic-mode",
type=bool,
default=False,
required=False,
help="Whether to use agentic mode.",
)
parser.add_argument(
"--agentic-max-concurrency",
type=int,
default=4,
required=False,
help="Max concurrency for agentic mode.",
)
parser.add_argument(
"--agentic-off-policy-steps",
type=int,
default=0,
required=False,
help="Number of off-policy steps for agentic mode.",
)

# Parse arguments
args = parser.parse_args()

if args.agentic_mode:
from tunix.rl.experimental import agentic_grpo_learner as grpo_learner # pylint: disable=g-import-not-at-top
else:
from tunix.rl import grpo_learner # pylint: disable=g-import-not-at-top


def validata_args():
if args.data_source == "tfds":
Expand Down Expand Up @@ -1100,12 +1124,22 @@ def get_rollout_config(engine: str) -> base_rollout.RolloutConfig:
rollout_config=get_rollout_config(args.rollout_engine),
)

grpo_config = grpo_learner.GRPOConfig(
num_generations=NUM_GENERATIONS,
num_iterations=NUM_ITERATIONS,
beta=BETA,
epsilon=EPSILON,
)
if args.agentic_mode:
grpo_config = grpo_learner.GRPOConfig(
num_generations=NUM_GENERATIONS,
num_iterations=NUM_ITERATIONS,
beta=BETA,
epsilon=EPSILON,
max_concurrency=args.agentic_max_concurrency,
off_policy_steps=args.agentic_off_policy_steps,
)
else:
grpo_config = grpo_learner.GRPOConfig(
num_generations=NUM_GENERATIONS,
num_iterations=NUM_ITERATIONS,
beta=BETA,
epsilon=EPSILON,
)


# RL cluster
Expand Down Expand Up @@ -1161,7 +1195,7 @@ def get_rollout_config(engine: str) -> base_rollout.RolloutConfig:

show_hbm_usage("Right before training")
with training_mesh:
grpo_trainer.train(train_dataset, eval_ds=val_dataset)
grpo_trainer.train(train_dataset, val_dataset)

# Load checkpoint first.

Expand Down
Loading