FSDP with non-uniform 'requires_grad'

abdelkrim_zitouni: <p>Hello, I’m currently trying to wrap my Model which contains some frozen params with nested FSDP. I read from the <code>pytorch.distributed.fsdp</code> this warning:</p>
<blockquote>
<p>FSDP has some constraints on freezing parameters (i.e. setting <code>param.requires_grad=False</code> ). For <code>use_orig_params=False</code> , each FSDP instance must manage parameters that are all frozen or all non-frozen. For <code>use_orig_params=True</code> , FSDP supports mixing frozen and non-frozen, but we recommend not doing so since then the gradient memory usage will be higher than expected (namely, equivalent to not freezing those parameters). This means that ideally, frozen parameters should be isolated into their own <code>nn.Module</code> s and wrapped separately with FSDP.</p>
</blockquote>
<p>I didn’t want to set <code>use_orig_params=True</code> and as they stated, I Isolated the frozen parameters into their own <code>nn.Module</code> s and then wrapped every <strong>submodule</strong> with FSDP, but I’m still getting this error:</p>
<p>ValueError: <code>FlatParameter</code> requires uniform <code>requires_grad</code></p>
<p>Could someone please clarify what they meant by ‘ideally’ or help me overcome this issue?</p>
<p>Thanks!</p>

cmitch29: <p>Hi, did you ever figure this out? I am trying to do something similar.</p>

abdelkrim_zitouni: <p>Hi, sorry for the late reply. I overcame the issue by using a type-based auto-wrap policy, which lets you exclude the modules with frozen parameters in FSDP <code>auto_wrapper_callable</code>:</p>
<pre data-code-wrap="py"><code class="lang-py"># Automatic wrapping sub-modules with inner FSDP
    auto_wrap_policy = None
    auto_wrapper_callable = None
    if FLAGS.auto_wrap_policy != "none":
        if FLAGS.auto_wrap_policy == "size_based":
            # auto-wrap all sub-modules with a minimum number of parameters (default 1e6)
            auto_wrap_policy = partial(
                size_based_auto_wrap_policy,
                min_num_params=int(float(FLAGS.auto_wrap_min_num_params)))
        elif FLAGS.auto_wrap_policy == "type_based":
            # auto-wrap all sub-modules in MBConv or FFN
            auto_wrap_policy = partial(
                transformer_auto_wrap_policy,
                transformer_layer_cls={
                    Backbone.MBConv, # you can specify the modules you want to include here
                    Backbone.FeedForward,
                })
        else:
            raise Exception(
                f"Invalid auto-wrap policy: {FLAGS.auto_wrap_policy}")
        if FLAGS.use_gradient_checkpointing:
            # Apply gradient checkpointing to auto-wrapped sub-modules if specified
            auto_wrapper_callable = lambda m, *args, **kwargs: FSDP(
                checkpoint_module(m), *args, **kwargs)

    def fsdp_wrap(m): return FSDP(
        m,
        compute_dtype=getattr(torch, FLAGS.compute_dtype),
        fp32_reduce_scatter=FLAGS.fp32_reduce_scatter,
        flatten_parameters=FLAGS.flatten_parameters,
        pin_layout_in_collective_ops=FLAGS.pin_layout_in_collective_ops,
        auto_wrap_policy=auto_wrap_policy,
        auto_wrapper_callable=auto_wrapper_callable)
</code></pre>
<p>I don’t know if this helps in your case.</p>