[FSDP][torch.compile] accelerator.unwrap_model and trainer._save work incorrectly when FSDP + torch.compile

[efsotr]

System Info

transformers 4.51.3
accelerate 1.6.0

Who can help?

@zach-huggingface @SunMarc

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

To use torch.compile, you need to either uninstall the kernels library or set the environment variable DISABLE_KERNEL_MAPPING to 1.

train.py

from typing import cast

import torch
from transformers import HfArgumentParser, Trainer, TrainingArguments, LlamaForCausalLM, LlamaConfig

args = HfArgumentParser(TrainingArguments)
training_args = cast(TrainingArguments, args.parse_args_into_dataclasses())[0]
print(training_args, flush=True)

config = LlamaConfig(
    vocab_size=128, 
    hidden_size=128, 
    intermediate_size=128*2,
    num_hidden_layers=2
)
model = LlamaForCausalLM(config).cuda().bfloat16()

train_dataset = [{"input_ids": torch.randint(0, 128, (128,)),
                  "labels": torch.randint(0, 128, (128,))} for i in range(16)]

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=train_dataset,
)

trainer.train()
trainer.save_state()

fsdp.yaml

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
fsdp_config:
    fsdp_sharding_strategy: FULL_SHARD
    fsdp_activation_checkpointing: false
    fsdp_use_orig_params: true
    fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
    fsdp_backward_prefetch_policy: BACKWARD_PRE
    fsdp_offload_params: false
    fsdp_state_dict_type: FULL_STATE_DICT
    fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer,Embedding
mixed_precision: 'no'
enable_cpu_affinity: false
machine_rank: 0
main_training_function: main
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

launch script

export CUDA_VISIBLE_DEVICES=0,1
export DISABLE_KERNEL_MAPPING=1

OUTPUT_DIR=test_fsdp
mkdir -p $OUTPUT_DIR

OMP_NUM_THREADS=8 accelerate launch --main_process_port 40129 --config_file fsdp.yaml \
     train.py \
    --torch_compile_mode default \
    --do_train \
    --optim adamw_torch_fused \
    --learning_rate 1e-3 \
    --weight_decay 0 \
    --lr_scheduler_type constant_with_warmup \
    --warmup_ratio 0.1 \
    --per_device_train_batch_size 4 \
    --per_device_eval_batch_size 4 \
    --eval_on_start 0 \
    --eval_strategy epoch \
    --eval_steps 1 \
    --save_strategy epoch \
    --save_only_model 1 \
    --greater_is_better False \
    --logging_strategy steps \
    --logging_steps 1 \
    --include_tokens_per_second \
    --output_dir $OUTPUT_DIR \
    --num_train_epochs 1 \
    --seed 0 \
    --report_to none \
    > $OUTPUT_DIR/training.log 2>&1 

Expected behavior

file test_fsdp/checkpoint-2/config.json exists

run

from safetensors import safe_open
path = "test_fsdp/checkpoint-2/model.safetensors"
file = safe_open(path, framework="pt")
print(file.keys())
lm_head = "lm_head.weight"
if lm_head not in file.keys():
    lm_head += "_orig_mod."
print(file.get_tensor(lm_head).shape)

expected to get

['lm_head.weight', 'model.embed_tokens.weight', 'model.layers.0.input_layernorm.weight', 'model.layers.0.mlp.down_proj.weight', 'model.layers.0.mlp.gate_proj.weight', 'model.layers.0.mlp.up_proj.weight', 'model.layers.0.post_attention_layernorm.weight', 'model.layers.0.self_attn.k_proj.weight', 'model.layers.0.self_attn.o_proj.weight', 'model.layers.0.self_attn.q_proj.weight', 'model.layers.0.self_attn.v_proj.weight', 'model.layers.1.input_layernorm.weight', 'model.layers.1.mlp.down_proj.weight', 'model.layers.1.mlp.gate_proj.weight', 'model.layers.1.mlp.up_proj.weight', 'model.layers.1.post_attention_layernorm.weight', 'model.layers.1.self_attn.k_proj.weight', 'model.layers.1.self_attn.o_proj.weight', 'model.layers.1.self_attn.q_proj.weight', 'model.layers.1.self_attn.v_proj.weight', 'model.norm.weight']
torch.Size([128, 128])

instead of

['_orig_mod.lm_head.weight', '_orig_mod.model.embed_tokens.weight', '_orig_mod.model.layers.0.input_layernorm.weight', '_orig_mod.model.layers.0.mlp.down_proj.weight', '_orig_mod.model.layers.0.mlp.gate_proj.weight', '_orig_mod.model.layers.0.mlp.up_proj.weight', '_orig_mod.model.layers.0.post_attention_layernorm.weight', '_orig_mod.model.layers.0.self_attn.k_proj.weight', '_orig_mod.model.layers.0.self_attn.o_proj.weight', '_orig_mod.model.layers.0.self_attn.q_proj.weight', '_orig_mod.model.layers.0.self_attn.v_proj.weight', '_orig_mod.model.layers.1.input_layernorm.weight', '_orig_mod.model.layers.1.mlp.down_proj.weight', '_orig_mod.model.layers.1.mlp.gate_proj.weight', '_orig_mod.model.layers.1.mlp.up_proj.weight', '_orig_mod.model.layers.1.post_attention_layernorm.weight', '_orig_mod.model.layers.1.self_attn.k_proj.weight', '_orig_mod.model.layers.1.self_attn.o_proj.weight', '_orig_mod.model.layers.1.self_attn.q_proj.weight', '_orig_mod.model.layers.1.self_attn.v_proj.weight', '_orig_mod.model.norm.weight']
torch.Size([8128])

If the --eval_strategy epoch in the launch script is changed to --eval_strategy no, then

['_orig_mod.lm_head.weight', '_orig_mod.model.embed_tokens.weight', '_orig_mod.model.layers.0.input_layernorm.weight', '_orig_mod.model.layers.0.mlp.down_proj.weight', '_orig_mod.model.layers.0.mlp.gate_proj.weight', '_orig_mod.model.layers.0.mlp.up_proj.weight', '_orig_mod.model.layers.0.post_attention_layernorm.weight', '_orig_mod.model.layers.0.self_attn.k_proj.weight', '_orig_mod.model.layers.0.self_attn.o_proj.weight', '_orig_mod.model.layers.0.self_attn.q_proj.weight', '_orig_mod.model.layers.0.self_attn.v_proj.weight', '_orig_mod.model.layers.1.input_layernorm.weight', '_orig_mod.model.layers.1.mlp.down_proj.weight', '_orig_mod.model.layers.1.mlp.gate_proj.weight', '_orig_mod.model.layers.1.mlp.up_proj.weight', '_orig_mod.model.layers.1.post_attention_layernorm.weight', '_orig_mod.model.layers.1.self_attn.k_proj.weight', '_orig_mod.model.layers.1.self_attn.o_proj.weight', '_orig_mod.model.layers.1.self_attn.q_proj.weight', '_orig_mod.model.layers.1.self_attn.v_proj.weight', '_orig_mod.model.norm.weight']
torch.Size([128, 128])

[SunMarc]

Thanks for the report @efsotr, can you share what do you have in checkpoint-2 folder. We save the model twice for each checkpoint (one that we use if we want to retrain the model from this specific checkpoint, one that is compatible with from_pretrained). See _save_checkpoint -> save_model and _save_optimizer_and_scheduler.

[efsotr]

in checkpoint-2 folder, model.safetensors, trainer_state.json, training_args.bin

[efsotr]

@SunMarc I set save_only_model to 1, so only save_model is called.

[SunMarc]

Then in _save_checkpoint, can you check that save_model is called and we executing save_pretrained afterwards and not torch.save ?

    def _save_checkpoint(self, model, trial):
        # In all cases, including ddp/dp/deepspeed, self.model is always a reference to the model we
        # want to save except FullyShardedDDP.
        # assert unwrap_model(model) is self.model, "internal model should be a reference to self.model"

        # Save model checkpoint
        checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"

        if self.hp_search_backend is None and trial is None:
            self.store_flos()

        run_dir = self._get_output_dir(trial=trial)
        output_dir = os.path.join(run_dir, checkpoint_folder)
        self.save_model(output_dir, _internal_call=True)

In save_model

        elif self.is_fsdp_enabled:
            if ("FULL_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type)) and (
                version.parse(accelerate_version) > version.parse("0.24.1")
            ):
                state_dict = self.accelerator.get_state_dict(self.model)
                if self.args.should_save:
                    self._save(output_dir, state_dict=state_dict)

In _save

    def _save(self, output_dir: Optional[str] = None, state_dict=None):
        # If we are executing this function, we are the process zero, so we don't check for that.
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"Saving model checkpoint to {output_dir}")

        supported_classes = (PreTrainedModel,) if not is_peft_available() else (PreTrainedModel, PeftModel)
        # Save a trained model and configuration using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        if not isinstance(self.model, supported_classes):
            if state_dict is None:
                state_dict = self.model.state_dict()

            if isinstance(self.accelerator.unwrap_model(self.model), supported_classes):
                self.accelerator.unwrap_model(self.model).save_pretrained(
                    output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
                )
            else:
                logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
                if self.args.save_safetensors:
                    safetensors.torch.save_file(
                        state_dict, os.path.join(output_dir, SAFE_WEIGHTS_NAME), metadata={"format": "pt"}
                    )
                else:
                    torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
        else:
            self.model.save_pretrained(
                output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
            )

[efsotr]

"Trainer.model is not a PreTrainedModel, only saving its state dict." in the log

[SunMarc]

then something might be broken when we try to unwrap the model. Can you check that this gives you ?
self.accelerator.unwrap_model(self.model) and try to find why the unwrap is not proceding as expected ?

in unwrap_model -> extract_model_from_parallel, we have the following:

    if is_torch_version(">=", FSDP_PYTORCH_VERSION) and is_torch_distributed_available():
        from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP

        options += (FSDP,)

    while isinstance(model, options):
        model = model.module

cc @S1ro1

[efsotr]

arg keep_torch_compile in accelerator.unwrap_model is True

[SunMarc]

Oh ... 🫠 Thanks for finding this 😅 Would you like to submit a PR to fix this ? we should remove torch.compile wrapper when saving the model in general.

[Joaquinecc]

Hey @SunMarc, I am new to open source. If no one is working on this, i would love to contribute it.

Best regards,
Joaquin

[Joaquinecc]

I wasn’t able to reproduce the same output when I ran with --torch_compile_mode default—TorchDynamo immediately threw an error. After removing that flag, the script completed successfully: it generated config.json and produced an unwrapped model with the following parameters:

[
  "lm_head.weight",
  "model.embed_tokens.weight",
  "model.layers.0.input_layernorm.weight",
  "model.layers.0.mlp.down_proj.weight",
  "model.layers.0.mlp.gate_proj.weight",
  "model.layers.0.mlp.up_proj.weight",
  "model.layers.0.post_attention_layernorm.weight",
  "model.layers.0.self_attn.k_proj.weight",
  "model.layers.0.self_attn.o_proj.weight",
  "model.layers.0.self_attn.q_proj.weight",
  "model.layers.0.self_attn.v_proj.weight",
  "model.layers.1.input_layernorm.weight",
  "model.layers.1.mlp.down_proj.weight",
  "model.layers.1.mlp.gate_proj.weight",
  "model.layers.1.mlp.up_proj.weight",
  "model.layers.1.post_attention_layernorm.weight",
  "model.layers.1.self_attn.k_proj.weight",
  "model.layers.1.self_attn.o_proj.weight",
  "model.layers.1.self_attn.q_proj.weight",
  "model.layers.1.self_attn.v_proj.weight",
  "model.norm.weight"
]

[SunMarc]

What kind of error id you have. I think that in general, we should remove unwrap torch.compile when saving the model. Maybe you can do a simpler example when you wrap a model with torch.compile and you try to save it to reproduce the error. Then, fix all occurrences where we try to save unwrapped model.