Skip to content

fix checkpoint loading for unsharded target encoder - #2729

Draft
Jubeku wants to merge 1 commit into
develop-ssl-diffusion-v1from
jk/develop-ssl-diffusion-v1/load-unsharded-weights
Draft

fix checkpoint loading for unsharded target encoder#2729
Jubeku wants to merge 1 commit into
develop-ssl-diffusion-v1from
jk/develop-ssl-diffusion-v1/load-unsharded-weights

Conversation

@Jubeku

@Jubeku Jubeku commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Description

Issue Number

Is this PR a draft? Mark it as draft.

Checklist before asking for review

  • I have performed a self-review of my code
  • My changes comply with basic sanity checks:
    • I have fixed formatting issues with ./scripts/actions.sh lint
    • I have run unit tests with ./scripts/actions.sh unit-test
    • I have documented my code and I have updated the docstrings.
    • I have added unit tests, if relevant
  • I have tried my changes with data and code:
    • I have run the integration tests with ./scripts/actions.sh integration-test
    • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
    • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
  • I have informed and aligned with people impacted by my change:
    • for config changes: the MatterMost channels and/or a design doc
    • for changes of dependencies: the MatterMost software development channel


is_model_sharded = cf.with_ddp and cf.with_fsdp
# Determine shardedness from the model instance itself, NOT from the global config flags:
is_model_sharded = any(type(p) is torch.distributed.tensor.DTensor for p in model.parameters())

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

It would be better to then have it as an argument to load_model()

@github-actions github-actions Bot added the model Related to model training or definition (not generic infra) label Aug 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

model Related to model training or definition (not generic infra)

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

2 participants