Skip to content

[hparams] refactor: split training_args.py into per-algorithm package#163

Merged
Jayce-Ping merged 3 commits into
mainfrom
refactor/split-training-args-into-package
May 24, 2026
Merged

[hparams] refactor: split training_args.py into per-algorithm package#163
Jayce-Ping merged 3 commits into
mainfrom
refactor/split-training-args-into-package

Conversation

@Jayce-Ping

Copy link
Copy Markdown
Collaborator

Summary

  • Convert the 989-line monolithic training_args.py into a package (training_args/) with one file per algorithm, improving isolation for parallel development and reducing merge conflicts as the algorithm count grows
  • Zero breaking changes: all existing import paths (from .training_args import X) resolve identically via the package __init__.py
  • Update ff-new-algorithm skill and sample_lifecycle.md docs to reflect the new file paths

Package Structure

src/flow_factory/hparams/training_args/
├── __init__.py   # re-exports (preserves all existing import paths)
├── _base.py      # EvaluationArguments, TrainingArguments, shared utilities
├── _registry.py  # _TRAINING_ARGS_REGISTRY + get_training_args_class()
├── grpo.py       # GRPOTrainingArguments
├── nft.py        # NFTTrainingArguments
├── awm.py        # AWMTrainingArguments
├── dpo.py        # DPOTrainingArguments
├── dgpo.py       # DGPOTrainingArguments (extends GRPO)
└── crd.py        # CRDTrainingArguments

Test plan

  • python -c "from flow_factory.hparams import GRPOTrainingArguments, get_training_args_class; print('OK')"
  • python -c "from flow_factory.hparams import get_training_args_class; cls = get_training_args_class('grpo'); print(cls.__name__)"
  • Verify existing YAML configs load correctly via Arguments.load_from_yaml()
  • grep -rn "from.*training_args" src/ shows no broken imports
  • Run existing test suite

🤖 Generated with Claude Code

Convert the 989-line monolithic training_args.py into a package with
one file per algorithm, improving isolation for parallel development
and reducing merge conflicts as the algorithm count grows.

Structure:
  training_args/
    __init__.py   - re-exports (preserves all existing import paths)
    _base.py      - EvaluationArguments, TrainingArguments, utilities
    _registry.py  - registry dict + get_training_args_class()
    grpo.py       - GRPOTrainingArguments
    nft.py        - NFTTrainingArguments
    awm.py        - AWMTrainingArguments
    dpo.py        - DPOTrainingArguments
    dgpo.py       - DGPOTrainingArguments (extends GRPO)
    crd.py        - CRDTrainingArguments

Zero breaking changes: `from .training_args import X` resolves
identically whether training_args is a .py file or a package __init__.

Also updates ff-new-algorithm skill and sample_lifecycle docs to
reflect the new file paths.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Copilot AI review requested due to automatic review settings May 24, 2026 01:20

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR refactors flow_factory.hparams.training_args from a single monolithic module into a package with per-algorithm modules, while preserving the existing public import surface via training_args/__init__.py.

Changes:

  • Split training_args.py into training_args/ with _base.py (shared args/utilities), _registry.py (lookup/registry), and one file per algorithm.
  • Preserve existing import paths by re-exporting all public symbols from training_args/__init__.py.
  • Update agent skill and lifecycle docs to reflect the new file locations.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

Show a summary per file
File Description
src/flow_factory/hparams/training_args/_base.py Introduces shared EvaluationArguments/TrainingArguments and shared standardization utilities.
src/flow_factory/hparams/training_args/_registry.py Adds centralized registry and dynamic lookup for algorithm-specific TrainingArguments subclasses.
src/flow_factory/hparams/training_args/__init__.py Re-exports classes/functions to keep existing imports working.
src/flow_factory/hparams/training_args/grpo.py Adds GRPOTrainingArguments in its own module.
src/flow_factory/hparams/training_args/nft.py Adds NFTTrainingArguments in its own module.
src/flow_factory/hparams/training_args/awm.py Adds AWMTrainingArguments in its own module.
src/flow_factory/hparams/training_args/dpo.py Adds DPOTrainingArguments in its own module.
src/flow_factory/hparams/training_args/dgpo.py Adds DGPOTrainingArguments (extends GRPO) in its own module.
src/flow_factory/hparams/training_args/crd.py Adds CRDTrainingArguments in its own module.
src/flow_factory/hparams/training_args.py Removes the old monolithic module after splitting.
.agents/skills/ff-new-algorithm/SKILL.md Updates the “new algorithm” agent skill to match the new package structure and registration steps.
.agents/knowledge/topics/sample_lifecycle.md Updates docs to reference TrainingArguments in training_args/_base.py.
Comments suppressed due to low confidence (1)

src/flow_factory/hparams/training_args/_base.py:278

  • Same issue as EvaluationArguments: the width override warning is emitted but self.resolution is never updated, so width= cannot override the resolution and self.width will not reflect the user-provided value.
        if self.width is not None and self.resolution[1] != self.width:
                logger.warning(
                    f"Both `resolution={self.resolution}` and `width={self.width}` are set. "
                    f"Using width to override: ({self.resolution[0]}, {self.width})."
                )

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +84 to +88
if self.width is not None and self.resolution[1] != self.width:
logger.warning(
f"Both `resolution={self.resolution}` and `width={self.width}` are set. "
f"Using width to override: ({self.resolution[0]}, {self.width})."
)
Comment on lines +45 to +54
def get_training_args_class(identifier: str) -> Type[TrainingArguments]:
"""
Resolve the TrainingArguments subclass for a given trainer type.

Supports:
1. Registry lookup: 'grpo' -> GRPOTrainingArguments
2. Direct python path: 'my_package.hparams.CustomTrainingArgs' -> CustomTrainingArgs

Falls back to base TrainingArguments if lookup fails.
"""
- Fix pre-existing bug where `width` override logged a warning but never
  actually updated `self.resolution` (in both EvaluationArguments and
  TrainingArguments __post_init__)
- Fix misleading docstring in get_training_args_class() that claimed
  "falls back to base TrainingArguments" when it actually raises ImportError

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (1)

src/flow_factory/hparams/training_args/_base.py:276

  • The height/width override block in TrainingArguments.__post_init__ is over-indented (extra indentation level), which is inconsistent with the surrounding code and will be reformatted by Black. Adjust indentation to a single level under each if.
        if self.height is not None and self.resolution[0] != self.height:
                logger.warning(
                    f"Both `resolution={self.resolution}` and `height={self.height}` are set. "
                    f"Using height to override: ({self.height}, {self.resolution[1]})."
                )
                self.resolution = (self.height, self.resolution[1])
        if self.width is not None and self.resolution[1] != self.width:
                logger.warning(

Comment on lines +78 to +85
if self.height is not None and self.resolution[0] != self.height:
logger.warning(
f"Both `resolution={self.resolution}` and `height={self.height}` are set. "
f"Using height to override: ({self.height}, {self.resolution[1]})."
)
self.resolution = (self.height, self.resolution[1])
if self.width is not None and self.resolution[1] != self.width:
logger.warning(
Comment on lines +292 to +293
world_size = get_world_size()
logger.info("World Size:" + str(world_size))
- Fix over-indented height/width override blocks in both
  EvaluationArguments and TrainingArguments (8 extra spaces → standard)
- Replace string concatenation with f-string in world_size log
- Add blank line before EvaluationArguments.__post_init__ for consistency

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

@Jayce-Ping Jayce-Ping merged commit 3274a0f into main May 24, 2026
3 of 4 checks passed
87003697 pushed a commit to 87003697/Flow-Factory that referenced this pull request Jun 2, 2026
Brings trellis2 branch up-to-date with origin/main, integrating:
- PR X-GenGroup#170: DiffusionOPD on-policy distillation trainer
- PR X-GenGroup#168: Multi-dataset training with per-source reward routing
- PR X-GenGroup#165: GenEval reward + unified trainer sampling pipeline
- PR X-GenGroup#163: hparams/training_args.py → package split
- PR X-GenGroup#162: Reconcile config with runtime distributed state
- PRs X-GenGroup#121,X-GenGroup#146-X-GenGroup#161: CRD algorithm, Docker CUDA 12.9, HF resume, etc.

Conflict resolutions:
- trainers/registry.py: merged both sides (trellis2_grpo/nft + crd/opd)
- rewards/registry.py: merged both sides (trellis2 rewards + geneval)
- hparams/training_args.py: deleted (accept main's package split), added
  trellis2_grpo/nft entries to _registry.py
- trainers/grpo.py: removed duplicate evaluate() (unified in BaseTrainer)
- trainers/nft.py: adopted generate_samples(), removed duplicate evaluate()
- trainers/abc.py: added _extra_eval_inference_kwargs() hook to BaseTrainer
  evaluate() so Trellis2TrainerMixin can inject stages/render_kwargs
- samples/samples.py: merged source/source_id fields with trellis2 additions
- rewards/reward_processor.py: merged source-aware gating with trellis2's
  _store_reward_extra_info
- data_utils/loader.py: merged multi-source pipeline with image_3d dataset
- data_utils/sampler_loader.py: merged refactored parameters
- hparams/args.py: merged multi-source alignment with trellis2's group-aligned
- guidance/rewards.md: merged both (PickScore_TextImage_Sum + GenEval)

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants