[FSDP] LoRA tensors seem to vanish into oblivion when using FullyShardedDataParallel and peft

Oleg_Silkin: <p>Hello all,</p>
<p>We recently started using FSDP through the <img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=12" title=":hugs:" class="emoji" alt=":hugs:" loading="lazy" width="20" height="20"> Accelerate library and are running into weird issues when trying to train with LoRA from the <img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=12" title=":hugs:" class="emoji" alt=":hugs:" loading="lazy" width="20" height="20"> peft library.</p>
<p>I’ll describe the current issue I’m facing and will also discuss a few other things that I’ve tried doing.</p>
<p><strong>Current issue</strong></p>
<p>With my current setup, we can <em>run</em> the forward pass of an FSDP model with LoRA, but we <em>cannot save</em> the model.</p>
<p>Saving LoRA checkpoints seems to fail as the LoRA layers seemingly lose their dimensionality by the time that we call the function to save the checkpoint. To illustrate what I mean, please see the following error:</p>
<pre><code class="lang-auto">[rank0]: Traceback (most recent call last):
[rank0]:   File "/home/ec2-user/training/src/instructlab/training/main_ds.py", line 805, in &lt;module&gt;
[rank0]:     main(args)
[rank0]:   File "/home/ec2-user/training/src/instructlab/training/main_ds.py", line 613, in main
[rank0]:     train(
[rank0]:   File "/home/ec2-user/training/src/instructlab/training/main_ds.py", line 462, in train
[rank0]:     save_hf_format_accelerate(
[rank0]:   File "/home/ec2-user/training/src/instructlab/training/utils.py", line 692, in save_hf_format_accelerate
[rank0]:     model.module.merge_adapter()
[rank0]:   File "/home/ec2-user/training/venv/lib64/python3.11/site-packages/peft/tuners/tuners_utils.py", line 550, in merge_adapter
[rank0]:     module.merge(adapter_names=adapter_names)
[rank0]:   File "/home/ec2-user/training/venv/lib64/python3.11/site-packages/peft/tuners/lora/layer.py", line 481, in merge
[rank0]:     delta_weight = self.get_delta_weight(active_adapter)
[rank0]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/home/ec2-user/training/venv/lib64/python3.11/site-packages/peft/tuners/lora/layer.py", line 548, in get_delta_weight
[rank0]:     output_tensor = transpose(weight_B @ weight_A, self.fan_in_fan_out) * self.scaling[adapter]
[rank0]:                               ~~~~~~~~~^~~~~~~~~~
[rank0]: RuntimeError: inconsistent tensor size, expected tensor [4096] and src [11008] to have the same number of elements, but got 4096 and 11008 elements respectively
</code></pre>
<p>Which is a really strange error considering the context in how these tensors are initialized (<a href="https://github.com/huggingface/peft/blob/aa3bd8fbf6820bb000e148a289b6b0c48c339774/src/peft/tuners/lora/layer.py#L121" rel="noopener nofollow ugc">from <img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=12" title=":hugs:" class="emoji" alt=":hugs:" loading="lazy" width="20" height="20"> peft</a>:</p>
<pre data-code-wrap="py"><code class="lang-py">    def update_layer(
        self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights, use_rslora, use_dora: bool = False
    ):

        # ...

        # Actual trainable parameters
        self.lora_A[adapter_name] = nn.Linear(self.in_features, r, bias=False)
        self.lora_B[adapter_name] = nn.Linear(r, self.out_features, bias=False)
</code></pre>
<p>Inspecting the LoRA adapters at this stage via <code>distributed.breakpoint()</code> produces the following results:</p>
<pre data-code-wrap="py"><code class="lang-py">&gt; self.lora_A[adapter_name].weight.shape
(11006, 4)
&gt; self.lora_B[adapter_name].weight.shape
(4, 4096)
</code></pre>
<p>But in the <code>self.get_delta_weight</code> method where it attempts to perform the matrix multiplication between these two Tensors, inspecting their shapes yields the following result:</p>
<pre data-code-wrap="py"><code class="lang-py">&gt; self.lora_A[active_adapter].weight.shape
(11006,)
&gt; self.lora_B[active_adapter].weight.shape
(4096,)
</code></pre>
<p>Which seems to imply that <code>(fan_in + fan_out) x (lora_rank - 1)</code> parameters simply disappeared at some point. I tried searching for where this possibly could be happening but was unsuccessful.</p>
<p>Here’s my current setup:</p>
<p><code>PeftConfig</code>:</p>
<pre data-code-wrap="py"><code class="lang-py">def get_fsdp_config(args, model):
    # Third Party
    from accelerate.utils import FullyShardedDataParallelPlugin
    from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload
    from peft.utils.other import fsdp_auto_wrap_policy

    fsdp_plugin = FullyShardedDataParallelPlugin(
        auto_wrap_policy=fsdp_auto_wrap_policy(model),
        limit_all_gathers=True,
        mixed_precision_policy=MixedPrecision(
            param_dtype=torch.bfloat16,
            reduce_dtype=torch.bfloat16,
            buffer_dtype=torch.bfloat16,
        ),
        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
        sharding_strategy=ShardingStrategy[args.fsdp_sharding_strategy],
        cpu_offload=CPUOffload(args.cpu_offload_params_fsdp),
        use_orig_params=True,
        cpu_ram_efficient_loading=True,
        sync_module_states=False,
    )
    return fsdp_plugin


peft_config = LoraConfig(
    lora_alpha=args.lora_alpha,
    lora_dropout=args.lora_dropout,
    r=args.lora_r,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=args.lora_target_modules,
)
model = get_peft_model(model, peft_config)
accel_args = {
    "fsdp_plugin": get_fsdp_config(args, model),
}
accelerator = Accelerator(**accel_args, mixed_precision="bf16")
accelerator.even_batches = False
return accelerator
</code></pre>
<h2><a name="p-456344-other-things-ive-tried-1" class="anchor" href="#p-456344-other-things-ive-tried-1"></a><strong>Other Things I’ve tried</strong></h2>
<h3><a name="p-456344-use_orig_paramsfalse-leads-to-an-error-about-data-not-allocated-yet-2" class="anchor" href="#p-456344-use_orig_paramsfalse-leads-to-an-error-about-data-not-allocated-yet-2"></a><code>use_orig_params=False</code> leads to an error about data not allocated yet</h3>
<p>With the current mixed precision settings and <code>use_orig_params=False</code> we cannot run training as for some reason we get a different about the data for the layers not being allocated yet:</p>
<pre><code class="lang-auto">[rank0]:   File "/home/ec2-user/training/venv/lib64/python3.11/site-packages/peft/tuners/lora/layer.py", line 548, in get_delta_weight
[rank0]:     output_tensor = transpose(weight_B @ weight_A, self.fan_in_fan_out) * self.scaling[adapter]
[rank0]:                               ~~~~~~~~~^~~~~~~~~~
[rank0]: RuntimeError: The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.
</code></pre>
<h3><a name="p-456344-use_orig_paramsfalse-and-mixed_precisionnone-leads-to-data-type-mismatch-between-float32-and-bfloat16-3" class="anchor" href="#p-456344-use_orig_paramsfalse-and-mixed_precisionnone-leads-to-data-type-mismatch-between-float32-and-bfloat16-3"></a><code>use_orig_params=False</code> and <code>mixed_precision=None</code> leads to data type mismatch between float32 and bfloat16</h3>
<p>When we have this configuration, we run into a problem where 3 iterations into the forward pass, we have the <code>inputs.dtype = float32</code> being passed into a Linear layer with <code>weight.dtype = bfloat16</code>:</p>
<pre><code class="lang-auto">[rank1]:   File "/home/ec2-user/training/venv/lib64/python3.11/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/ec2-user/training/venv/lib64/python3.11/site-packages/torch/nn/modules/linear.py", line 117, in forward
[rank1]:     return F.linear(input, self.weight, self.bias)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::BFloat16
</code></pre>
<h3><a name="p-456344-lora-mixed-precision-error-4" class="anchor" href="#p-456344-lora-mixed-precision-error-4"></a><strong>LoRA mixed precision error</strong></h3>
<p>When <code>use_orig_params=True</code> and no mixed precision setting is given to FSDP, then we again run into the data type mismatch on the very first forward pass:</p>
<pre><code class="lang-auto">[rank1]:   File "/home/ec2-user/training/venv/lib64/python3.11/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 858, in forward
[rank1]:     output = self._fsdp_wrapped_module(*args, **kwargs)
[rank1]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/ec2-user/training/venv/lib64/python3.11/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/ec2-user/training/venv/lib64/python3.11/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/ec2-user/training/venv/lib64/python3.11/site-packages/torch/nn/modules/linear.py", line 117, in forward
[rank1]:     return F.linear(input, self.weight, self.bias)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::BFloat16
</code></pre>
<p>Does anyone here have an idea for why this is happening? Any insight would be greatly appreciated!</p>

pratikkorat26: <p>Can you show me your code that you are using for saving your model ?</p>

Oleg_Silkin: <p><a class="mention" href="/u/pratikkorat26">@pratikkorat26</a> Sure, here’s the function that saves checkpoints:</p>
<pre data-code-wrap="py"><code class="lang-py">def save_hf_format_accelerate(
    args,
    model,
    tokenizer,
    accelerator: Accelerator,
    samples_seen,
    convert_granite=True,
    is_lora=False,
):
    log_rank_0(
        f"\033[93mSaving model in huggingface format at samples_seen: {samples_seen}\033[0m",
        to_print=True,
    )
    start = time.time()

    final_output_dir = Path(args.output_dir) / "hf_format" / f"samples_{samples_seen}"
    if args.is_granite and convert_granite:
        tmpdir = TemporaryDirectory("w")  # pylint: disable=consider-using-with
        output_dir = Path(tmpdir.name)
    else:
        output_dir = final_output_dir

    CONFIG_NAME = "config.json"
    output_config_file = output_dir / CONFIG_NAME

    get_state_dict_unpatched = accelerator.get_state_dict

    def _get_state_dict_patched(model, unwrap=False):
        return get_state_dict_unpatched(model, unwrap=unwrap)

    accelerator.get_state_dict = _get_state_dict_patched

    if accelerator.is_main_process:
        if is_lora:
            model.module.merge_adapter()
            model_state = model.module.state_dict()

        output_dir.mkdir(parents=True, exist_ok=True)
        if not model.module.config.architectures and convert_granite:
            model.module.config.architectures = ["LlamaForCausalLM"]
            warnings.warn(
                f"Adding architectures to ckpt: {model.module.config.architectures}",
            )
        model.module.config.to_json_file(output_config_file)
        tokenizer.save_pretrained(output_dir)

        if is_lora:
            save_dict_accelerate(
                accelerator,
                model_state,
                save_directory=output_dir,
                max_shard_size="5GB",
                safe_serialization=True,
            )
            model.module.unmerge_adapter()

    if not is_lora:
        accelerator.save_model(
            model,
            save_directory=output_dir,
            max_shard_size="5GB",
            safe_serialization=True,
        )

    if args.is_granite and convert_granite and accelerator.is_main_process:
        # export doesnt like the directory to exist
        if final_output_dir.exists():
            shutil.rmtree(final_output_dir)
        export_to_huggingface(
            pretrained_model_name_or_path=tmpdir.name,
            save_path=final_output_dir,
            model_type="llama",
        )
        tmpdir.cleanup()

    log_rank_0(f"\033[93mModel saved in {final_output_dir}\033[0m", to_print=True)
    log_rank_0(f"saving took {time.time() - start} seconds")
    dist.barrier()

    accelerator.get_state_dict = get_state_dict_unpatched
</code></pre>
<p><code>granite</code> here just translates to padding-free transformers. The save logic here was written for ZeRO Stage-2 when we were using DeepSpeed, and the FSDP code I’ve been running has been with <code>SHARD_GRAD_OP</code> as the sharding strategy. I’ve also tried setting the state dict type to <code>FULL_STATE_DICT</code> with no success.</p>