class_docstr: 实现对于 pytorch 自己实现的 fully sharded data parallel；请阅读
  `该文档 <https://pytorch.org/docs/stable/fsdp.html#torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict>`_
  了解更多：
  
  .. note::
  
      ``TorchFSDPDriver`` 大部分行为与 ``TorchDDPDriver`` 相同，如果您不了解 ``TorchDDPDriver``，
      您可以先阅读 :class:`~fastNLP.core.drivers.TorchDDPDriver`；
  
  .. warning::
  
      ``TorchFSDPDriver`` 现在还不支持断点重训功能，但是支持保存模型和加载模型；
  
      注意当您在加载和保存模型的 checkpointcallback 的时候，您可以通过在初始化 ``Trainer`` 时传入
      ``torch_kwargs={"fsdp_kwargs": {'save_on_rank0': True/False, 'load_on_rank0': True/False}}`` 来指定保存模型的行为：
  
          1. save/load_on_rank0 = True：表示在加载和保存模型时将所有 rank 上的模型参数全部聚合到 rank0 上，注意这样可能会造成 OOM；
          2. save/load_on_rank0 = False：表示每个 rank 分别保存加载自己独有的模型参数；
  
  :param model: 传入给 ``Trainer`` 的 ``model`` 参数
  :param parallel_device: 用于分布式训练的 ``gpu`` 设备
  :param is_pull_by_torch_run: 标志当前的脚本的启动是否由 ``python -m torch.distributed.launch`` 启动的
  :param fp16: 是否开启 fp16 训练
  :param torch_kwargs: 
  
      * *fsdp_kwargs* --
      * *set_grad_to_none* -- 是否在训练过程中在每一次 optimizer 更新后将 grad 置为 ``None``
      * *non_blocking* -- 表示用于 :meth:`torch.Tensor.to` 方法的参数 non_blocking
      * *gradscaler_kwargs* -- 用于 ``fp16=True`` 时，提供给 :class:`torch.amp.cuda.GradScaler` 的参数
  :kwargs:
      * *wo_auto_param_call* (``bool``) -- 是否关闭在训练时调用我们的 ``auto_param_call`` 函数来自动匹配 batch 和前向函数的参数的行为
  
      .. note::
  
          关于该参数的详细说明，请参见 :class:`~fastNLP.core.controllers.Trainer` 中的描述；函数 ``auto_param_call`` 详见 :func:`fastNLP.core.utils.auto_param_call`。
class_docstr_tok: ["\u5b9e\u73b0\u5bf9\u4e8e","pytorch","\u81ea\u5df1\u5b9e\u73b0\u7684","fully","sharded","data","parallel\uff1b\u8bf7\u9605\u8bfb","`","\u8be5\u6587\u6863","<","https",":","\/\/pytorch.org\/docs\/stable\/fsdp.html","#","torch.distributed.fsdp.FullyShardedDataParallel.full_optim_state_dict",">","`","_","\u4e86\u89e3\u66f4\u591a\uff1a","..","note",":",":","``","TorchFSDPDriver","``","\u5927\u90e8\u5206\u884c\u4e3a\u4e0e","``","TorchDDPDriver","``","\u76f8\u540c\uff0c\u5982\u679c\u60a8\u4e0d\u4e86\u89e3","``","TorchDDPDriver","``","\uff0c","\u60a8\u53ef\u4ee5\u5148\u9605\u8bfb",":","class",":","`","~fastNLP.core.drivers.TorchDDPDriver","`","\uff1b","..","warning",":",":","``","TorchFSDPDriver","``","\u73b0\u5728\u8fd8\u4e0d\u652f\u6301\u65ad\u70b9\u91cd\u8bad\u529f\u80fd\uff0c\u4f46\u662f\u652f\u6301\u4fdd\u5b58\u6a21\u578b\u548c\u52a0\u8f7d\u6a21\u578b\uff1b","\u6ce8\u610f\u5f53\u60a8\u5728\u52a0\u8f7d\u548c\u4fdd\u5b58\u6a21\u578b\u7684","checkpointcallback","\u7684\u65f6\u5019\uff0c\u60a8\u53ef\u4ee5\u901a\u8fc7\u5728\u521d\u59cb\u5316","``","Trainer","``","\u65f6\u4f20\u5165","``","torch_kwargs=","{","``","fsdp_kwargs","''",":","{","'save_on_rank0","'",":","True\/False",",","'load_on_rank0","'",":","True\/False","}","}","``","\u6765\u6307\u5b9a\u4fdd\u5b58\u6a21\u578b\u7684\u884c\u4e3a\uff1a","1.","save\/load_on_rank0","=","True\uff1a\u8868\u793a\u5728\u52a0\u8f7d\u548c\u4fdd\u5b58\u6a21\u578b\u65f6\u5c06\u6240\u6709","rank","\u4e0a\u7684\u6a21\u578b\u53c2\u6570\u5168\u90e8\u805a\u5408\u5230","rank0","\u4e0a\uff0c\u6ce8\u610f\u8fd9\u6837\u53ef\u80fd\u4f1a\u9020\u6210","OOM\uff1b","2.","save\/load_on_rank0","=","False\uff1a\u8868\u793a\u6bcf\u4e2a","rank","\u5206\u522b\u4fdd\u5b58\u52a0\u8f7d\u81ea\u5df1\u72ec\u6709\u7684\u6a21\u578b\u53c2\u6570\uff1b",":","param","model",":","\u4f20\u5165\u7ed9","``","Trainer","``","\u7684","``","model","``","\u53c2\u6570",":","param","parallel_device",":","\u7528\u4e8e\u5206\u5e03\u5f0f\u8bad\u7ec3\u7684","``","gpu","``","\u8bbe\u5907",":","param","is_pull_by_torch_run",":","\u6807\u5fd7\u5f53\u524d\u7684\u811a\u672c\u7684\u542f\u52a8\u662f\u5426\u7531","``","python","-m","torch.distributed.launch","``","\u542f\u52a8\u7684",":","param","fp16",":","\u662f\u5426\u5f00\u542f","fp16","\u8bad\u7ec3",":","param","torch_kwargs",":","*","*","fsdp_kwargs","*","--","*","*","set_grad_to_none","*","--","\u662f\u5426\u5728\u8bad\u7ec3\u8fc7\u7a0b\u4e2d\u5728\u6bcf\u4e00\u6b21","optimizer","\u66f4\u65b0\u540e\u5c06","grad","\u7f6e\u4e3a","``","None","``","*","*","non_blocking","*","--","\u8868\u793a\u7528\u4e8e",":","meth",":","`","torch.Tensor.to","`","\u65b9\u6cd5\u7684\u53c2\u6570","non_blocking","*","*","gradscaler_kwargs","*","--","\u7528\u4e8e","``","fp16=True","``","\u65f6\uff0c\u63d0\u4f9b\u7ed9",":","class",":","`","torch.amp.cuda.GradScaler","`","\u7684\u53c2\u6570",":","kwargs",":","*","*","wo_auto_param_call","*","(","``","bool","``",")","--","\u662f\u5426\u5173\u95ed\u5728\u8bad\u7ec3\u65f6\u8c03\u7528\u6211\u4eec\u7684","``","auto_param_call","``","\u51fd\u6570\u6765\u81ea\u52a8\u5339\u914d","batch","\u548c\u524d\u5411\u51fd\u6570\u7684\u53c2\u6570\u7684\u884c\u4e3a","..","note",":",":","\u5173\u4e8e\u8be5\u53c2\u6570\u7684\u8be6\u7ec6\u8bf4\u660e\uff0c\u8bf7\u53c2\u89c1",":","class",":","`","~fastNLP.core.controllers.Trainer","`","\u4e2d\u7684\u63cf\u8ff0\uff1b\u51fd\u6570","``","auto_param_call","``","\u8be6\u89c1",":","func",":","`","fastNLP.core.utils.auto_param_call","`","\u3002"]
docstr: 从 ``filepath`` 中加载权重并赋值到当前 driver 的模型上。
  
  :param filepath: 加载权重或模型的路径
  :param load_state_dict: 保存的内容是否只是权重；在 ``TorchFSDPDriver`` 中只能为 ``True`` 。
  :param kwargs:
  :return:
docstr_tok: ["\u4ece","``","filepath","``","\u4e2d\u52a0\u8f7d\u6743\u91cd\u5e76\u8d4b\u503c\u5230\u5f53\u524d","driver","\u7684\u6a21\u578b\u4e0a\u3002",":","param","filepath",":","\u52a0\u8f7d\u6743\u91cd\u6216\u6a21\u578b\u7684\u8def\u5f84",":","param","load_state_dict",":","\u4fdd\u5b58\u7684\u5185\u5bb9\u662f\u5426\u53ea\u662f\u6743\u91cd\uff1b\u5728","``","TorchFSDPDriver","``","\u4e2d\u53ea\u80fd\u4e3a","``","True","``","\u3002",":","param","kwargs",":",":","return",":"]
code: def load_model(
      self,
      filepath: Union[Path, str],
      only_state_dict: bool = True,
      **kwargs,
  ):
      """
      从 ``filepath`` 中加载权重并赋值到当前 driver 的模型上。
  
      :param filepath: 加载权重或模型的路径
      :param load_state_dict: 保存的内容是否只是权重；在 ``TorchFSDPDriver`` 中只能为 ``True`` 。
      :param kwargs:
      :return:
      """
      if only_state_dict is False:
          raise RuntimeError(
              "When using `TorchFSDPDriver`, only `only_state_dict=True` is allowed."
          )
      filepath = Path(filepath)
      prefix = filepath.parent
      filename = filepath.name
      _filename = filename.split(".")
      filename, suffix = _filename[0], ".".join(_filename[1:])
  
      if not self._load_on_rank0:
          _filename = filename.split("_")
          filename = (
              _filename[0]
              + f"_rank{int(os.environ.get(FASTNLP_GLOBAL_RANK, 0))}_"
              + _filename[1]
          )
          filepath = prefix.joinpath(filename + "." + suffix)
          states = torch.load(filepath)
      else:
          states = torch.load(filepath, map_location="cpu")
  
      if isinstance(states, dict) and only_state_dict is False:
          logger.rank_zero_warning(
              f"It seems like that {filepath} only contains state, you may need to use "
              f"`only_state_dict=True`"
          )
      elif not isinstance(states, dict) and only_state_dict is True:
          logger.rank_zero_warning(
              f"It seems like that {filepath} is not state, you may need to use "
              f"`only_state_dict=False`"
          )
      if not isinstance(states, Mapping):
          states = states.state_dict()
  
      if self._load_on_rank0:
          with FullyShardedDataParallel.state_dict_type(
              self.model, StateDictType.FULL_STATE_DICT
          ):
              self.model.load_state_dict(states)
      else:
          with FullyShardedDataParallel.state_dict_type(
              self.model, StateDictType.LOCAL_STATE_DICT
          ):
              self.model.load_state_dict(states)
  
imports: from fastNLP.envs.imports import _TORCH_GREATER_EQUAL_1_12, _NEED_IMPORT_TORCH
  import os
  from typing import Optional, Union, List, Dict, Mapping
  from pathlib import Path
  from .ddp import TorchDDPDriver
  from fastNLP.core.drivers.torch_driver.utils import _DDPWrappingModel
  from fastNLP.envs import FASTNLP_DISTRIBUTED_CHECK, FASTNLP_MODEL_FILENAME, FASTNLP_CHECKPOINT_FILENAME, FASTNLP_GLOBAL_RANK, rank_zero_call
  from fastNLP.core.drivers.torch_driver.utils import DummyGradScaler
  from fastNLP.core.log import logger
  from fastNLP.core.utils import check_user_specific_params
  from .utils import optimizer_state_to_device
label_desc: Use image node_id 5 for calling the TorchFSDPDriver obj's underlying member method code with example usage: obj.load_model(filepath, only_state_dict) without return types