class_docstr: Strategy for distributed training using torch.distributed.
class_docstr_tok: ["Strategy","for","distributed","training","using","torch.distributed","."]
code: def save_pretrained(
      self,
      model: nn.Module,
      path: str,
      shard: bool = False,
      tokenizer: Optional[PreTrainedTokenizerBase] = None,
  ) -> None:
      if dist.get_rank() == 0:
          unwrapped_model = self.unwrap_model(model)
          assert isinstance(
              unwrapped_model, (Actor, Critic, RewardModel)
          )
          pretrained_model = unwrapped_model.model
          assert isinstance(pretrained_model, PreTrainedModel)
          # HACK: only use hf save_pretrained to save config
          pretrained_model.save_pretrained(
              path, save_function=lambda *args, **kwargs: None
          )
          if tokenizer is not None:
              tokenizer.save_pretrained(path)
  
      model_path = os.path.join(path, "pytorch_model.bin")
      self.save_model(model, model_path, shard=shard)
  
      def _replace_keys(model_path: str, replace_fn: Callable):
          state_dict = torch.load(model_path, map_location="cpu")
          state_dict = {replace_fn(k): v for k, v in state_dict.items()}
          torch.save(state_dict, model_path)
  
      # FIXME: save_model would add "model." prefix to keys of pytorch_model.bin
      # HACK: rename keys of pytorch_model.bin
      if dist.get_rank() == 0:
          _replace_keys(
              model_path, lambda k: k.replace("model.", "", 1)
          )
  
imports: import os
  import random
  from collections import OrderedDict
  from typing import Callable, Optional
  import numpy
  import torch
  import torch.distributed
  import torch.nn
  from coati.experience_buffer import ExperienceBuffer
  from coati.models import Actor, Critic, RewardModel
  from torch.utils.data import DataLoader
  from transformers.modeling_utils import PreTrainedModel
  from transformers.tokenization_utils_base import PreTrainedTokenizerBase
  from colossalai.booster.plugin import TorchDDPPlugin
  from colossalai.booster.plugin.torch_ddp_plugin import TorchDDPModel
  from .base import Strategy
  from .sampler import DistributedSampler
label_desc: Use image node_id 9 for calling the DDPStrategy obj's underlying member method code with example usage: obj.save_pretrained(model, path, shard, tokenizer) without return types