class_docstr:     Callback for valuate the performance of the model.
  Args:
      actor_num_params: The number of parameters of the actor model.
      critic_num_params: The number of parameters of the critic model.
      initial_model_num_params: The number of parameters of the initial model.
      reward_model_num_params: The number of parameters of the reward model.
      enable_grad_checkpoint: Whether to enable gradient checkpointing.
      ignore_episodes: The number of episodes to ignore when calculating the performance.
class_docstr_tok: ["Callback","for","valuate","the","performance","of","the","model",".","Args",":","actor_num_params",":","The","number","of","parameters","of","the","actor","model",".","critic_num_params",":","The","number","of","parameters","of","the","critic","model",".","initial_model_num_params",":","The","number","of","parameters","of","the","initial","model",".","reward_model_num_params",":","The","number","of","parameters","of","the","reward","model",".","enable_grad_checkpoint",":","Whether","to","enable","gradient","checkpointing",".","ignore_episodes",":","The","number","of","episodes","to","ignore","when","calculating","the","performance","."]
code: def on_fit_end(self) -> None:
      avg_duration = all_reduce_mean(
          self.timer.duration, self.coordinator.world_size
      )
      avg_throughput = (
          self.num_samples * self.dp_world_size / (avg_duration + 1e-12)
      )
      mp_world_size = self.coordinator.world_size // self.dp_world_size
      avg_tflops_per_gpu_megatron = (
          self.flop_megatron
          / 1e12
          / (avg_duration + 1e-12)
          / mp_world_size
      )
      avg_tflops_per_gpu = (
          self.flop / 1e12 / (avg_duration + 1e-12) / mp_world_size
      )
      self.coordinator.print_on_master(
          f"num_samples: {self.num_samples}, dp_world_size: {self.dp_world_size}, flop_megatron: {self.flop_megatron}, flop: {self.flop}, avg_duration: {avg_duration}, "
          f"avg_throughput: {avg_throughput}"
      )
      self.coordinator.print_on_master(
          f"Throughput: {avg_throughput:.2f} samples/sec, TFLOPS per GPU by Megatron: {avg_tflops_per_gpu_megatron:.2f}, TFLOPS per GPU: {avg_tflops_per_gpu:.2f}"
      )
  
file_path: ColossalAI/examples/language/llama2/performance_evaluator.py
imports: from time import time
  from typing import Optional
  import torch
  import torch.distributed
  from torch import Tensor
  import colossalai.utils.device
  from colossalai.cluster import DistCoordinator
  from colossalai.utils.device import get_current_device
label_desc: Use image node_id 4 for calling the PerformanceEvaluator obj's underlying member method code with example usage: obj.on_fit_end() without return types