code: def _forward_backward_test_impl(
      self,
      forward_only: bool,
      fwd_bwd_func: FwdStepFunc,
      pipeline_model_parallel_world_size: Optional[int],
      virtual_pipeline_model_parallel_size: Optional[int],
      async_comm: bool = False,
      *,
      default_backend: Optional[str] = None,
      p2p_backend: Optional[str] = None,
  ) -> None:
      if fwd_bwd_func == _forward_backward_pipelining_with_interleaving:
          self.assertIsNotNone(virtual_pipeline_model_parallel_size)
          self.assertGreater(virtual_pipeline_model_parallel_size, 1)
      dtype_options = (
          self.dtypes
          or [torch.float32, torch.double] + _get_autocast_dtypes()
      )
  
      for dtype, deallocate_pipeline_outputs in itertools.product(
          dtype_options,
          self.deallocate_options,
      ):
          grad_scaler = (
              torch.cuda.amp.GradScaler(init_scale=4.0)
              if dtype == torch.half
              else None
          )
  
          (
              tensor_model_parallel_world_size,
              data_parallel_size,
              pipeline_model_parallel_world_size,
          ) = _get_default_world_sizes_model_parallel_world_size(
              pipeline_model_parallel_world_size
          )
  
          parallel_state.initialize_model_parallel(
              tensor_model_parallel_size_=tensor_model_parallel_world_size,
              pipeline_model_parallel_size_=pipeline_model_parallel_world_size,
              virtual_pipeline_model_parallel_size_=virtual_pipeline_model_parallel_size,
              default_backend=default_backend,
              p2p_backend=p2p_backend,
          )
          pp_utils._reconfigure_microbatch_calculator(
              rank=parallel_state.get_tensor_model_parallel_rank(),
              rampup_batch_size=None,
              global_batch_size=self.GLOBAL_BATCH_SIZE,
              micro_batch_size=self.MICRO_BATCH_SIZE,
              data_parallel_size=parallel_state.get_data_parallel_world_size(),
          )
  
          global_batch_shape = (
              self.GLOBAL_BATCH_SIZE
              // parallel_state.get_data_parallel_world_size(),
              self.HIDDEN_SIZE,
              self.HIDDEN_SIZE,
          )
  
          batch = None
          if parallel_state.is_pipeline_first_stage():
              batch = (
                  torch.ones(global_batch_shape, dtype=dtype).cuda(),
              )
  
          model = build_model(
              testing_utils.model_provider_func,
              # Use DDP only when it's better to have
              wrap_with_ddp=data_parallel_size > 1,
              virtual_pipeline_model_parallel_size=virtual_pipeline_model_parallel_size,
              hidden_size=self.HIDDEN_SIZE,
          )
  
          offset = (
              pipeline_model_parallel_world_size
              if virtual_pipeline_model_parallel_size is not None
              else 0
          )
          for idx, model_module in enumerate(model):
              model_module = model_module.to(dtype)
  
          _param_groups = _get_params_for_weight_decay_optimization(
              model
          )
          optimizer = torch.optim.Adam(_param_groups, lr=1e-3)
  
          pp_utils.update_num_microbatches(0)
  
          for _ in range(self.NUM_FWD_BWD_ITERATIONS):
              loss = fwd_bwd_func(
                  testing_utils.fwd_step_func,
                  batch,
                  model,
                  forward_only=forward_only,
                  # `tensor_shape` is the shape of micro batch.
                  tensor_shape=(
                      self.MICRO_BATCH_SIZE,
                      self.HIDDEN_SIZE,
                      self.HIDDEN_SIZE,
                  ),
                  dtype=dtype,
                  async_comm=async_comm,
                  grad_scaler=grad_scaler,
                  deallocate_pipeline_output=deallocate_pipeline_outputs,
              )
  
          parallel_state.destroy_model_parallel()
  
imports: import os
  import logging
  import itertools
  from typing import Optional, Tuple, List
  import unittest
  import torch
  from torch.testing._internal import common_utils
  from torch.testing._internal import common_cuda
  from torch.testing._internal import common_distributed
  from apex._autocast_utils import _get_autocast_dtypes
  from apex.transformer import parallel_state
  from apex.transformer.pipeline_parallel import utils
  from apex.transformer.pipeline_parallel.schedules.common import FwdStepFunc, build_model, _get_params_for_weight_decay_optimization
  from apex.transformer.pipeline_parallel.schedules.fwd_bwd_no_pipelining import forward_backward_no_pipelining
  from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_with_interleaving import _forward_backward_pipelining_with_interleaving
  from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_without_interleaving import forward_backward_pipelining_without_interleaving
  from apex.transformer.testing.distributed_test_base import UccDistributedTestBase
  from apex.transformer.testing import commons
label_desc: Use image node_id 3 for calling the UccPipelineParallelForwardBackwardProf obj's underlying member method code with example usage: obj._forward_backward_test_impl(forward_only, fwd_bwd_func, pipeline_model_parallel_world_size, virtual_pipeline_model_parallel_size, async_comm) without return types