huggingface / huggingface/transformers

[resume optimization] skip loading pretrained weights on resume

Open
#11,465 6 comments 0 reactions 1 assignee Claimed by @stas00 View on GitHub
WIP
Dominant language
Python
Stars
166k
Forks
34.6k
Avg merge
3d 8h
Merged PRs (30d)
276

Description

This is similar to what was discussed in https://github.com/huggingface/transformers/issues/9205, which proposed not to random init weights on `from_pretrained`, but this time it's about resume - currently we load pretrained weights and immediately drop them on resume from checkpoint in Trainer.

To solve this we, for example, could change examples:

1. to figure out the checkpoint immediately after we init `TrainingArguments` and just before model is created.
2. then change `from_pretrained()` API to do keep everything as is, except loading the weights from `state_dict`, if say `skip_weights_load=True` is passed:

So the code becomes:
```

if training_args.do_train:
if last_checkpoint is not None:
checkpoint = last_checkpoint
elif os.path.isdir(model_args.model_name_or_path):
checkpoint = model_args.model_name_or_path
else:
checkpoint = None

model = AutoModelForSeq2SeqLM.from_pretrained(
model_args.model_name_or_path,
[...],
skip_weights_load=checkpoint is not None,
)

if training_args.do_train:
train_result = trainer.train(resume_from_checkpoint=checkpoint)
```

Any flaws in my thinking?

@patrickvonplaten, @sgugger

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.