[docs] add fsdp_tips.rst

[sshleifer]

This is a dumping ground to collect things we want to document about FSDP.

Please comment with anything random that comes to mind.

Once there have been no changes to FSDP for 48 hrs (or some other proxy for stability), I will finalize this and set it to Ready For Review.

[min-xu-ai]

should this include the new wrap/auto_wrap feature?

[sshleifer]

Yes leave desired markdown in comment else I can take a pass tomorrow.

[min-xu-ai]

#### The `enable_wrap` context

There are two cases where the `enable_wrap` context can be useful:

* When you'd like to apply the same parameters to all child modules that you wrap with FSDP. Calling the `wrap` function within the said context will save you from passing the same set of FSDP parameters explicitly.
* When wrapping large models that does NOT fit within the CPU memory. I.e. you don't first create the full model and then traverse it to wrap it with FSDP at different parts. Instead, you create a wrapped instance of the model incrementally as you build up the model, allowing large modules to be sharded in-place.

example:
        with enable_wrap(**fsdp_params):
            # Wraps layer in FSDP by default if within context
            self.l1 = wrap(torch.nn.Linear(5, 5))
            # Wraps children modules by default based on min_num_params
            self.l2 = auto_wrap(TransformerBlock(), min_num_params=1e8)

[min-xu-ai]

One more thing to be added to the doc:

If model weight initialization happens deterministically after sharding, the final weights will be comprised of N piece of identical shards. This would have negative effects on the model initial state.

@myleott @sshleifer Does the above looks good? If so, I can also add it to the doctoring of FSDP. We generate doc from the doctoring.

[sshleifer]

These are not at all perfect, please feel free to push changes to this branch or comment.
I made a separate fsdp_tips.rst and cross-linked with fsdp.rst to avoid one crowding out the other on the same page. There may be a nicer way to do this in rst, (e.g. linking to section 2 at the top of 1 document).

I also copied the contents of the docstring of _init_param_attributes into a section called "State management with extra parameter attributes", which is not shown in the other page cause it's private and useful to know.

[sshleifer]

fsdp_tips.html.zip. Unzip and open in put the path in the chrome URL bar.

[Vittorio-Caggiano]

just a small suggestion.
Would it possible to add also a small tutorial?

[myleott]

Would it possible to add also a small tutorial?

Yes, this is a great idea 😄

Just to cross-reference and not lose track, Min also suggested the tutorial should include custom weight init inside a summon_full_params context: #454 (comment)

[myleott]

This looks great, thanks @sshleifer! I made some comments below, but I think we can ship this and iterate. What do you think @min-xu-ai?

[myleott]

the new syntax is enable_wrap(wrapper_cls=FSDP, **fsdp_params)

[min-xu-ai]

I think we can ship this and iterate. What do you think @min-xu-ai?

Yeah, totally.