Skip to content

Refactor the code and rm accelerate dependency #1032

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

Jump to bottom
Merged
aoyulong merged 5 commits into from
Jan 7, 2026
Merged

Refactor the code and rm accelerate dependency #1032

aoyulong merged 5 commits into from
Jan 7, 2026

Conversation

@legitnull
Copy link
Collaborator

@legitnull legitnull commented Jan 5, 2026 •
edited
Loading

PR Category

Train

PR Types

Bug Fixes

PR Description

  • Add a flag use_amp to enable mixed precision training
  • Remove the accelerate dep, now the PI0/Pi0.5 training works with DDP
  • Remove the example pi0_old along with the model definition, train/infer/server script
  • Reorder the code structure
    • move transformation.diffusion to inference/core/diffusion
    • rename diffusion_entrypoint.py to inference_diffusion.py and move to inference/core

heavyrain-lzy
heavyrain-lzy reviewed Jan 6, 2026
View reviewed changes
@legitnull legitnull changed the title Fix PI0 training with AMP Remove the accelerate dependency Jan 6, 2026
Hchnr
Hchnr reviewed Jan 6, 2026
View reviewed changes
aoyulong
aoyulong previously approved these changes Jan 6, 2026
View reviewed changes
Copy link
Contributor

@aoyulong aoyulong left a comment

Choose a reason for hiding this comment

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

LGTM

@aoyulong aoyulong changed the title Remove the accelerate dependency Refactor the code and rm accelerate dependency Jan 7, 2026
@aoyulong aoyulong merged commit 3794714 into flagos-ai:main Jan 7, 2026
5 of 14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@heavyrain-lzy heavyrain-lzy heavyrain-lzy left review comments

@Hchnr Hchnr Hchnr left review comments

@aoyulong aoyulong aoyulong left review comments

@cyber-pioneer cyber-pioneer Awaiting requested review from cyber-pioneer cyber-pioneer is a code owner

@zhaoyinglia zhaoyinglia Awaiting requested review from zhaoyinglia zhaoyinglia is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@legitnull @aoyulong @heavyrain-lzy @Hchnr