Skip to content

Allow loading pretrained shared Pytorch checkpoints into flax models #18170

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

Closed
Sea-Snell wants to merge 3 commits into from
Closed

Allow loading pretrained shared Pytorch checkpoints into flax models #18170

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9b52d1d
load from sharded pytorch checkpoint for flax model
Sea-Snell Jul 18, 2022
06fc890
added import from utils
Sea-Snell Jul 21, 2022
ce558a8
ran make style
Sea-Snell Jul 21, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/transformers/modeling_flax_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
FLAX_WEIGHTS_INDEX_NAME,
FLAX_WEIGHTS_NAME,
HUGGINGFACE_CO_RESOLVE_ENDPOINT,
WEIGHTS_INDEX_NAME,
WEIGHTS_NAME,
EntryNotFoundError,
PushToHubMixin,
Expand Down Expand Up @@ -639,6 +640,10 @@ def from_pretrained(
if from_pt and os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)):
# Load from a PyTorch checkpoint
archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
elif from_pt and os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME)):
Copy link
Collaborator

@ArthurZucker ArthurZucker Jul 22, 2022

Choose a reason for hiding this comment

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

BTW this will only work if the WEIGHTS_INDEX_NAME file is locally present, and does not include the hub.

Copy link
Author

@Sea-Snell Sea-Snell Jul 24, 2022

Choose a reason for hiding this comment

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

Yeah, let's just finalize yours. What's left to do?

Copy link
Collaborator

@ArthurZucker ArthurZucker Jul 24, 2022

Choose a reason for hiding this comment

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

Maybe just fixing the tests, and making sure that the tests are actually good. Should be quiet straightforward☺️🙌

Copy link
Collaborator

@ArthurZucker ArthurZucker Jul 25, 2022

Choose a reason for hiding this comment

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

@Sea-Snell we just need to fix test_from_sharded_pt which is failing because the model used for comparison are not the same! Simply using the same model (either upload a new model using the same config but shard it with save_pretrained and setting the max_shard_size to 150KB should do the trick.

# Load from a sharded PyTorch checkpoint
archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME)
is_sharded = True
Comment on lines +643 to +646
Copy link
Collaborator

@ArthurZucker ArthurZucker Jul 22, 2022

Choose a reason for hiding this comment

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

LGTM, just wondering if could add a small test?
You can use hf-internal-testing/tiny-random-bert-sharded/.
Also I opened #18026 which is really similar, which adds

    @is_pt_flax_cross_test
    def test_from_sharded_pt(self):
        model = FlaxBertModel.from_pretrained("hf-internal-testing/tiny-random-bert-sharded", from_pt=True)
        ref_model = FlaxBertModel.from_pretrained("ArthurZ/tiny-random-bert-flax-only")
        for p1, p2 in zip(flatten_dict(model.params).values(), flatten_dict(ref_model.params).values()):
            assert np.allclose(np.array(p1), np.array(p2))

Was not really aware that the conversion would be straight forward let me have a look

elif os.path.isfile(os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)):
# Load from a Flax checkpoint
archive_file = os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)
Expand Down