class_docstr: Subclass of FSDPStrategy to ensure compatibility with Ray orchestration.
  
  For a full list of initialization arguments, please refer to:
  https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.strategies.FSDPStrategy.html
class_docstr_tok: ["Subclass","of","FSDPStrategy","to","ensure","compatibility","with","Ray","orchestration",".","For","a","full","list","of","initialization","arguments",",","please","refer","to",":","https",":","\/\/lightning.ai\/docs\/pytorch\/stable\/api\/lightning.pytorch.strategies.FSDPStrategy.html"]
docstr: Gathers the full state dict to rank 0 on CPU.
docstr_tok: ["Gathers","the","full","state","dict","to","rank","0","on","CPU","."]
code: def lightning_module_state_dict(self) -> Dict[str, Any]:
      """Gathers the full state dict to rank 0 on CPU."""
      assert (
          self.model is not None
      ), "Failed to get the state dict for a None model!"
  
      if _LIGHTNING_GREATER_EQUAL_2_0 and _TORCH_FSDP_AVAILABLE:
          with FullyShardedDataParallel.state_dict_type(
              module=self.model,
              state_dict_type=StateDictType.FULL_STATE_DICT,
              state_dict_config=FullStateDictConfig(
                  offload_to_cpu=True, rank0_only=True
              ),
          ):
              state_dict = self.model.state_dict()
              prefix_len = len("_forward_module.")
              return {k[prefix_len:]: v for k, v in state_dict.items()}
      else:
          # Otherwise Lightning uses Fairscale FSDP, no need to unshard by ourself.
          return super().lightning_module_state_dict()
  
imports: import logging
  import os
  import shutil
  import tempfile
  from typing import Any, Dict
  import torch
  from packaging.version import Version
  import ray
  from ray import train
  from ray._private.usage.usage_lib import TagKey, record_extra_usage_tag
  from ray.train import Checkpoint
  from ray.util import PublicAPI
label_desc: Use image node_id 4 for calling the RayFSDPStrategy obj's underlying member method code with example usage: obj.lightning_module_state_dict() and returns: super, unknown