Skip to content

Revamp/gather args infer#2

Merged
Rockdu merged 6 commits into
mainfrom
revamp/gather_args_infer
Jun 10, 2026
Merged

Revamp/gather args infer#2
Rockdu merged 6 commits into
mainfrom
revamp/gather_args_infer

Conversation

@Rockdu

@Rockdu Rockdu commented May 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

This pr:

  1. Move all implicit args inference logics into arguments.py for clear references (getattr(args, "xxx", default_val), (args.xxx or None), etc), so that almost all arguments and their inferences are determined after the argument parsing phase.
  2. update_weight_target_module was hard-coded but should be set in arguments.py, since it's related to a specific training run. This pr move update_weight_target_module into the arguments

@Rockdu
Rockdu requested review from Ying1123 and guapisolo as code owners May 13, 2026 08:18
Rockdu and others added 4 commits May 13, 2026 21:33
@Rockdu
remove debug prints
5b0ed6e
@Rockdu
remove unused diffusion offload functions and clean up imports in rol…
7fddb2c 
…lout.py
@Rockdu @claude
Merge origin/main into revamp/gather_args_infer
46798ad 
Resolve conflicts and move update_weight_target_module back to args:

- update_weight_target_module: removed from TrainPipelineConfig and
  registered as --update-weight-target-module in arguments.py; the
  weight updater now reads it from args. It is a training-side quantity,
  not a model property.
- actor.py: keep main's debug_train_only branch; updaters no longer
  receive an explicit target_module arg.
- placement_group.py: take main's simplified create_rollout_manager
  (num_gpus=0, plain pg); the PR's use_diffusion_rollout branch is dead
  code in our setup.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@Rockdu @claude
Restore main-only code the PR had deleted
a33e403 
Keep parity with upstream for code that exists on main:

- rollout.py: restore the [reward stats] stdout print in
  _convert_samples_to_train_data (raw + normalized mean/std/min/max).
  This is on main and is in the live rollout path; the PR had dropped it.
- loss.py: restore the custom_pg_loss_reducer selection so this file
  stays byte-identical to main. (loss.py is currently dead code in the
  diffusion repo, so this is parity-only; avoids future upstream-sync
  conflicts.)

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@Rockdu
Rockdu merged commit aedfb29 into main Jun 10, 2026
Rockdu added a commit to Rockdu/miles_diffusion that referenced this pull request Jun 10, 2026
@Rockdu @claude
chore: apply repo-wide pre-commit lint + revamp arg gather (radixark#2)
895bccf 
Squashed from the bootstrap branch's intermediate commits:
- Mass black / ruff / isort / autoflake fixes across miles/, flow_grpo/,
  train_diffusion.py to satisfy pre-commit hooks for the first time.
- Revamp gather-args inference path in miles/utils/arguments.py and
  related rollout/ray/router/backends call sites (originally PR radixark#2 on
  the fork).

This bundles WIP commits that interleaved CI bootstrap, docker image,
and code refactor. Squashed for a clean PR diff.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Rockdu added a commit to Rockdu/miles_diffusion that referenced this pull request Jun 10, 2026
@Rockdu @claude
chore: apply pre-commit lint to miles/ + revamp arg gather (radixark#2)
35516a8 
Squashed from the bootstrap branch's intermediate WIP commits:
- pre-commit (black / ruff / isort / autoflake) fixes across miles/ and
  train_diffusion.py
- Revamp gather-args inference path in miles/utils/arguments.py and
  related rollout/ray/router/backends call sites (originally PR radixark#2 on
  the fork; included here since the squash predates that PR landing on
  upstream/main from this branch's POV)

flow_grpo/ is reference-only and out of CI scope — explicitly NOT
touched (pyproject.toml already pins ruff extend-exclude flow_grpo).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
zhihengy pushed a commit that referenced this pull request Jun 30, 2026
@claude
fix(sp): round-0 review real issues -- production num_heads guard + s…
bb75b71 
…glang import guard

Two of the three real findings from the Codex review phase fixed immediately (the other is AC-9 full-RL, GPU-blocked):
- #3 [AC-2/3] the production num_heads%ulysses guard was bypassed: create_fsdp_parallel_state reads args.num_attention_heads (absent in production -> None -> skips the check). Added a late guard in apply_sequence_parallel using the real transformer.config.num_attention_heads, rejecting illegal ulysses at startup (torchrun verified: heads10%ulysses4 rejected)
- #2 guard tests/sglang_usp_import_guard.py: asserts the actually-imported sglang contains _AllToAllSingle/_RingFlashAttention + checksum covers dtype/shape, preventing env drift to an unpatched /sgl-workspace/sglang (2 passed)

Note: Codex gap#2 (sglang fixes missing) and blocking#1 (tests not runnable) verified as false positives -- Codex had not activated the miles-diffusion conda env: the real env imports the patched /workspace/.../sglang (editable), ray is inside the env (pytest 22 passed).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@guapisolo guapisolo Awaiting requested review from guapisolo guapisolo is a code owner
@Ying1123 Ying1123 Awaiting requested review from Ying1123

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@Rockdu