cc @muellerzr @SunMarc !

Thanks for the nice report ! This is indeed a very strange behavior. Could you try to see if you get the same model at the end with/without the callback. At first glance, it looks like with fsdp.FullyShardedDataParallel.summon_full_params(model) is the potential culprit. Could you try to just call alone in on_epoch_end ?

Hey @SunMarc! Just a note re:

Could you try to see if you get the same model at the end with/without the callback.

I am following this issue and also replied to the HuggingFace forum. TL;DR when unsharding the model, only the *.safetensors file differ between runs with and without the callback -- so we don't get the same model.

Hi @SunMarc! Thanks a lot for looking into this!

I confirm that I get different model weights depending on whether I use the callback or not. All three Safetensors files show up with diff.

I tried to run the generation in on_epoch_end without calling summon_full_params, but I get this error when I reach model.generate:

RuntimeError: 'weight' must be 2-D

I assume this behaviour is expected, based on this comment. I hope this is what you were asking, but do tell if I got it wrong.

I looked a bit more into PyTorch's documentation for summon_full_params, and tried setting writeback=False, just to make sure. But it has no effect: predictions after training look fine, but the checkpoint is wrong.

false

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

cc @XuehaiPan if you have some time !

    def on_epoch_end(
        self,
        args: transformers.TrainingArguments,
        state: transformers.TrainerState,
        control: transformers.TrainerControl,
        **kwargs,
    ) -> None:
        if state.epoch is None or int(state.epoch) % 25 != 0:
            return
        model = cast(transformers.PreTrainedModel, kwargs["model"])
        with torch.no_grad():
            self.run(model)

    @torch.no_grad()
    def run(self, model: transformers.PreTrainedModel) -> None:
        model.eval()

        for batch in self.dataset.iter(batch_size=7):
            encoding = self.tokenizer(batch["text"], return_tensors="pt", padding=True).to(model.device)

            with fsdp.FullyShardedDataParallel.summon_full_params(model):
                outputs = model.generate(
                    inputs=encoding.input_ids,
                    attention_mask=encoding.attention_mask,
                    pad_token_id=self.tokenizer.eos_token_id,
                    max_new_tokens=16,
                    do_sample=False,
                )

            predictions = self.tokenizer.batch_decode(
                outputs[:, encoding.input_ids.shape[1] :],  # Skip the returned prompt.
                skip_special_tokens=True,
                clean_up_tokenization_spaces=True,
            )

            if accelerate.PartialState().is_main_process:
                print(predictions)

Hi @pandrei7, I wonder have you ever tried to remove the model.eval() statement or change the model back to training mode after validation?

Hey @XuehaiPan, thanks! I tried both options:

1. removing model.eval() call
2. keeping model.eval() and adding model.train() at the end of the callback

Unfortunately, both options give corrupted checkpoints. Here's the inference on the dummy dataset from @pandrei7's example:

1. ['.20 0 0 20 20 0 ', '.0.20 0 0 20 0 0', '0 0 0 0 0 0 0 0', '0 0 0 0 0 0 0 0 ', '0 0 0 0 0 0 0 0 ']
2. ['times 0...... ', '0.......', 'equals.....0 equals. ', '.....0 equals..', '...0 equals....']

The correct inference result should be:

"Two"                 -> "times 10 equals 20."
"Two times"           -> "10 equals 20."
"Two times 10"        -> "equals 20."
"Two times 10 equals" -> "20."

Hi @XuehaiPan! Thanks for responding.

I confirm what @alexandru-dinu said: I tried both approaches, but the issue persists. Predictions are correct while training, but the checkpoint I load from disk generates incorrect completions.

I'm wondering if this might be an issue with the way I use PyTorch... Do you think I should try asking there as well?

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.