Skip to content

fix(scripts): fix checkpoint sharding restore, precompile CLI flag, and rollout sync in grpo demo #1962

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
Adityakk9031 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(scripts): fix checkpoint sharding restore, precompile CLI flag, and rollout sync in grpo demo #1962

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
29 changes: 22 additions & 7 deletions scripts/grpo_demo_llama3_qwen2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,12 +68,14 @@

print("This script is still WIP and try at your own discretion")

# Disable precompilation for faster iteration, need to toggle it back for
# official run
os.environ["SKIP_JAX_PRECOMPILE"] = "1"

# Parse command line options
parser = argparse.ArgumentParser(description="Arguments for GRPO demo")
parser.add_argument(
"--skip-jax-precompile",
action="store_true",
default=False,
help="Disable JAX precompilation for faster iteration.",
)
parser.add_argument(
"--root-dir",
type=str,
Expand Down Expand Up @@ -317,6 +319,11 @@
# Parse arguments
args = parser.parse_args()

if args.skip_jax_precompile:
os.environ["SKIP_JAX_PRECOMPILE"] = "1"
else:
os.environ.pop("SKIP_JAX_PRECOMPILE", None)


def validata_args():
if args.data_source == "tfds":
Expand Down Expand Up @@ -1285,9 +1292,12 @@ def get_rollout_config(engine: str) -> base_rollout.RolloutConfig:
)

filter_type = nnx.LoRAParam if ENABLE_LORA else nnx.Param
training_state = nnx.state(training_model, filter_type)
abs_params = jax.tree.map(
lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype),
nnx.state(training_model, filter_type),
lambda x: jax.ShapeDtypeStruct(
x.shape, x.dtype, sharding=getattr(x, "sharding", None)
),
training_state,
)
checkpointer = ocp.StandardCheckpointer()
trained_lora_params = checkpointer.restore(trained_ckpt_path, target=abs_params)
Expand All @@ -1296,11 +1306,16 @@ def get_rollout_config(engine: str) -> base_rollout.RolloutConfig:
training_model,
jax.tree.map(
lambda a, b: b,
nnx.state(training_model, filter_type),
training_state,
trained_lora_params,
),
)

# Synchronize restored weights to the rollout worker/sampler for post-training evaluation
filter_types = (filter_type,)
src_filtered_params = nnx.state(training_model, filter_types)
rl_engine.rollout.update_params(src_filtered_params, filter_types)

(eval_corr, eval_total, eval_accuracy, eval_partial_accuracy, eval_format_accuracy) = evaluate( # pylint: disable=unbalanced-tuple-unpacking
test_dataset,
rollout_sampler,
Expand Down
Loading