Composable FSDP Transform

[t-vi]

Make FSDP composable with other transforms.

Desired features:

• allow FSDPTransform to compose with other transforms (e.g. some Quantization or QLORA),
• allow FSDPTransform to work with Materialization:
  • sharding on meta,
  • sharding of state dict,
  • release memory after sharding if desired and called with a materialized model (in particular default for fsdp(jit(model))?), i.e. fix fsdp(jit(...)) transform can use more memory compared to jit(fsdp(...)) #478
• as a side condition fsdp(jit(model)) must continue to work as a convenience method for practitioners (thanks @crcrpar for the observation that this should stay simple).

Steps include:

• move transformation of module to FSDPTransform.transform_module (PR Move FSDP module transformation to transform_module #986 )
• do memory-releasing (PR free orig parameters after sharding #988 (also covering PR [fsdp as a transform] Free memories occupied by materialized params that are initialized with torch.device("meta") #932))
• implement FSDPTransform.transform_state_dict_for_submodule (PR implement FSDPTransform.transform_state_dict_for_submodule #989 )
• delegate materialization (PR update FSDPTransform to delegate materialization to the MaterializationTransform #995 )

For the side condition of keeping fsdp(jit(model)) working as is, all of these should work while keeping the existing tests unchanged.

@crcrpar @IvanYashchuk (who asked about the plan)

cc @carmocca @crcrpar