Deepspeed Zero3 resume training with Huggingface Trainer Failed

Hello,

When I try to resume training for a huggingface PEFT model, with the deepspeed launcher and the same number of gpus (8 gpus on the same node), I got lora size mismatch:

lora_weight = lora_B.weight @ lora_A.weight

  File "lib/python3.10/site-packages/torch/nn/modules/module.py", line 1561, in _call_impl
RuntimeError: size mismatch, got input (18944), mat (18944x8), vec (0)

Here's how I launch the training script

deepspeed src/train.py \
    ...
    --deepspeed_config zero3.json

Here's what's in the train.py:

training_args = TrainingArguments(
        output_dir=output_dir,
        overwrite_output_dir=True,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        dataloader_num_workers=1,
        num_train_epochs=3,

        gradient_accumulation_steps=16,
        save_total_limit=10,
        learning_rate=2e-4,
        weight_decay=0.,
        warmup_ratio=0.03,
        lr_scheduler_type="cosine",
        save_strategy="best",
        metric_for_best_model = "loss",
        # logging_steps=4,
        logging_steps=16,
        do_eval=True,
        eval_strategy='steps',
        eval_steps=1000,
        bf16=True,
        tf32=False,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        report_to=args.report_to,
        remove_unused_columns=False,
        deepspeed=args.deepspeed_config
    )

model and dataloader initialization etc, no deepspeed engine was called here...

    logging.info("Building trainer...")
    trainer = MultimodalLlamaTrainer(model=model_stack['model'],
                                     tokenizer=model_stack['tokenizer'],
                                     group_by_modality_length=True,
                                     args=training_args,
                                     **data_module)
    
    logging.info("Starting training...")
    if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
        trainer.train(resume_from_checkpoint=True)
    else:
        trainer.train()

Here's what's in my checkpoint dir:

When I use the script to train from scratch, it works perfectly. When I try to resume training, the model loaded successfully, but eventually got lora size mismatch. I can confirm that I did not change any configurations. How could I fix this?

Thanks a lot!

Start with the resume_from_checkpoint path in train.py and the referenced zero3.json configuration. Reproduce the failure using the provided DeepSpeed command and checkpoint setup, then compare resume behavior with training from scratch. Done means resuming the same PEFT training run completes without the reported LoRA size mismatch.