Skip to content

[Feature] Offload all text encoders by default #594

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
merged 10 commits into from
Jul 4, 2025
Merged

[Feature] Offload all text encoders by default #594

merged 10 commits into from
Jul 4, 2025

Conversation

Edenzzzz
Copy link
Collaborator

@Edenzzzz Edenzzzz commented Jul 3, 2025

Part of #572. Please check that this doesn't cause multi-node training to hang.
cc @SolitaryThinker

@SolitaryThinker SolitaryThinker added the go Trigger Buildkite CI label Jul 3, 2025
@SolitaryThinker SolitaryThinker self-requested a review July 3, 2025 04:31
@SolitaryThinker
Copy link
Collaborator

need to rebase for the buildkite CI fix

@Edenzzzz Edenzzzz force-pushed the offload_text_enc branch from 6c5f09f to 62ec225 Compare July 3, 2025 17:12
@Edenzzzz
Copy link
Collaborator Author

Edenzzzz commented Jul 3, 2025

@SolitaryThinker Done

@SolitaryThinker
Copy link
Collaborator

thanks, testing now on multi-node

SolitaryThinker
SolitaryThinker requested changes Jul 3, 2025
View reviewed changes
Copy link
Collaborator

@SolitaryThinker SolitaryThinker left a comment

Choose a reason for hiding this comment

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

where is the actual logic for offloading text encoders? I think you forgot to include that logic in this PR?

fastvideo/v1/models/loader/component_loader.py Outdated

# Load the module
return loader.load(component_model_path, architecture, fastvideo_args)
return loader.load(component_model_path, architecture, fastvideo_args)
Copy link
Collaborator

Choose a reason for hiding this comment

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

revert this?

fastvideo/v1/models/loader/fsdp_load.py Outdated

# choose `assign=True` since we cannot call `copy_` on meta tensor
return model.load_state_dict(sharded_sd, strict=strict, assign=True)
return model.load_state_dict(sharded_sd, strict=strict, assign=True)
Copy link
Collaborator

Choose a reason for hiding this comment

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

also here

fastvideo/v1/models/loader/weight_utils.py Outdated

# If there were no matches, return the untouched param name
return name
return name
Copy link
Collaborator

Choose a reason for hiding this comment

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

here

@Edenzzzz
Copy link
Collaborator Author

Edenzzzz commented Jul 3, 2025

where is the actual logic for offloading text encoders? I think you forgot to include that logic in this PR?

Added

@SolitaryThinker SolitaryThinker merged commit ed1e8d6 into hao-ai-lab:main Jul 4, 2025
1 check passed
@Edenzzzz Edenzzzz deleted the offload_text_enc branch July 4, 2025 00:46
@Edenzzzz Edenzzzz mentioned this pull request Jul 5, 2025
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@SolitaryThinker SolitaryThinker SolitaryThinker approved these changes

Assignees
No one assigned
Labels
go Trigger Buildkite CI
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@Edenzzzz @SolitaryThinker