class_docstr: 测试多卡情况下 save 和 load 相关函数的表现
class_docstr_tok: ["\u6d4b\u8bd5\u591a\u5361\u60c5\u51b5\u4e0b","save","\u548c","load","\u76f8\u5173\u51fd\u6570\u7684\u8868\u73b0"]
docstr: 测试save和load函数，主要测试 dataloader 被替换了 sampler 之后的情况
docstr_tok: ["\u6d4b\u8bd5save\u548cload\u51fd\u6570\uff0c\u4e3b\u8981\u6d4b\u8bd5","dataloader","\u88ab\u66ff\u6362\u4e86","sampler","\u4e4b\u540e\u7684\u60c5\u51b5"]
code: def test_save_and_load_with_bucketedbatchsampler(
      self, device, only_state_dict, fp16
  ):
      """
      测试save和load函数，主要测试 dataloader 被替换了 sampler 之后的情况
      """
  
      try:
          path = "model.ckp"
          num_replicas = len(device)
  
          dataloader = dataloader_with_bucketedbatchsampler(
              self.dataset,
              length=[10 for i in range(len(self.dataset))],
              batch_size=4,
              shuffle=True,
              drop_last=False,
          )
          dataloader.batch_sampler.set_distributed(
              num_replicas=int(os.getenv("WORLD_SIZE", "1")),
              rank=int(os.getenv("RANK", "0")),
              pad=True,
          )
          num_consumed_batches = 4
          driver1, driver2 = generate_driver(
              20,
              1,
              device=device,
              fp16=fp16,
              train_dataloader=dataloader,
          ), generate_driver(
              20,
              1,
              device=device,
              fp16=False,
              train_dataloader=dataloader,
          )
  
          already_seen_x_set = set()
          already_seen_y_set = set()
          driver1.set_sampler_epoch(dataloader, 4)
          for idx, batch in enumerate(dataloader):
              if idx >= num_consumed_batches:
                  break
              already_seen_x_set.update(
                  batch["x"]
                  .reshape(
                      -1,
                  )
                  .tolist()
              )
              already_seen_y_set.update(
                  batch["y"]
                  .reshape(
                      -1,
                  )
                  .tolist()
              )
  
          # 同步
          dist.barrier()
  
          # 保存状态
          sampler_states = dataloader.batch_sampler.state_dict()
          save_states = {"num_consumed_batches": num_consumed_batches}
          driver1.save_checkpoint(
              Path(path),
              save_states,
              dataloader,
              only_state_dict,
              should_save_model=True,
          )
          dist.barrier()
          # 加载
          # 更改 batch_size
          dataloader = dataloader_with_bucketedbatchsampler(
              self.dataset,
              length=[10 for i in range(len(self.dataset))],
              batch_size=2,
              shuffle=True,
              drop_last=False,
          )
          dataloader.batch_sampler.set_distributed(
              num_replicas=driver2.world_size,
              rank=driver2.global_rank,
              pad=True,
          )
          dist.barrier()
          load_states = driver2.load_checkpoint(
              Path(path),
              dataloader,
              only_state_dict,
              should_load_model=True,
          )
          dist.barrier()
          replaced_loader = load_states.pop("dataloader")
  
          # 1. 检查 optimizer 的状态
          # TODO optimizer 的 state_dict 总是为空
  
          # 2. 检查 batch_sampler 是否被正确地加载和替换
          assert not (replaced_loader is dataloader)
          assert (
              replaced_loader.batch_sampler is dataloader.batch_sampler
          )
          assert isinstance(
              replaced_loader.batch_sampler, BucketedBatchSampler
          )
          if os.environ["FASTNLP_GLOBAL_RANK"] == "0":
              assert (
                  replaced_loader.batch_sampler.seed
                  == sampler_states["seed"]
              )
          assert (
              replaced_loader.batch_sampler.num_consumed_samples
              == num_consumed_batches * 4 * num_replicas
          )
  
          # 4. 检查 model 的参数是否正确
          # 5. 检查 batch_idx
          start_batch = load_states.pop("batch_idx_in_epoch")
          assert start_batch == 2 * num_consumed_batches
          left_x_batches = set()
          left_y_batches = set()
          driver2.set_sampler_epoch(replaced_loader, 4)
          for idx, batch in enumerate(replaced_loader):
              left_x_batches.update(
                  batch["x"]
                  .reshape(
                      -1,
                  )
                  .tolist()
              )
              left_y_batches.update(
                  batch["y"]
                  .reshape(
                      -1,
                  )
                  .tolist()
              )
              batch = driver1.move_data_to_device(batch)
              res1 = driver1.model(
                  batch,
                  fastnlp_fn=driver1.model.module.model.evaluate_step,
                  # Driver.model -> DataParallel.module -> _FleetWrappingModel.model
                  fastnlp_signature_fn=None,
                  wo_auto_param_call=False,
              )
              res2 = driver2.model(
                  batch,
                  fastnlp_fn=driver2.model.module.model.evaluate_step,
                  fastnlp_signature_fn=None,
                  wo_auto_param_call=False,
              )
              assert torch.equal(res1["preds"], res2["preds"])
  
          assert (
              len(left_x_batches) + len(already_seen_x_set)
              == len(self.dataset) / num_replicas
          )
          assert (
              len(left_x_batches | already_seen_x_set)
              == len(self.dataset) / num_replicas
          )
          assert (
              len(left_y_batches) + len(already_seen_y_set)
              == len(self.dataset) / num_replicas
          )
          assert (
              len(left_y_batches | already_seen_y_set)
              == len(self.dataset) / num_replicas
          )
          dist.barrier()
      finally:
          rank_zero_rm(path)
  
imports: import os
  from pathlib import Path
  import pytest
  from fastNLP.core.drivers.torch_driver.deepspeed import DeepSpeedDriver
  from fastNLP.core.samplers import RandomSampler, BucketedBatchSampler, UnrepeatedRandomSampler
  from tests.helpers.models.torch_model import TorchNormalModel_Classification_1
  from tests.helpers.datasets.torch_data import TorchNormalXYDataset
  from tests.helpers.utils import magic_argv_env_context
  from fastNLP.envs.distributed import rank_zero_rm
  from fastNLP import logger
  from fastNLP.envs.imports import _NEED_IMPORT_TORCH, _NEED_IMPORT_DEEPSPEED
label_desc: Use image node_id 4 for calling the TestSaveLoad obj's underlying member method code with example usage: obj.test_save_and_load_with_bucketedbatchsampler(device, only_state_dict, fp16) without return types