code: def suite():
      suite = unittest.TestSuite()
  
      suite.addTest(GPT2FuncTestCase("test_mp1_gpu2_node1_fp16"))
  
      # Baseline = Megatron + Torch.Optim.Adam
      # Test = Megatron + Torch.Optim.Adam + ZeRO-Offload
      suite.addTest(
          GPT2FuncTestCase("test_mp1_gpu1_node1_zero2_torch_offload")
      )
      suite.addTest(
          GPT2FuncTestCase("test_mp1_gpu2_node1_zero2_torch_offload")
      )
      suite.addTest(
          GPT2FuncTestCase("test_mp2_gpu4_node1_zero2_torch_offload")
      )
      suite.addTest(
          GPT2FuncTestCase("test_mp4_gpu4_node1_zero2_torch_offload")
      )
  
      # Baseline = Megatron + Torch.Optim.Adam
      # Test = Megatron + DeepSpeedAdam + ZeRO-Offload
      suite.addTest(
          GPT2FuncTestCase("test_mp1_gpu1_node1_zero2_ds_offload")
      )
      suite.addTest(
          GPT2FuncTestCase("test_mp1_gpu2_node1_zero2_ds_offload")
      )
      suite.addTest(
          GPT2FuncTestCase("test_mp2_gpu4_node1_zero2_ds_offload")
      )
      suite.addTest(
          GPT2FuncTestCase("test_mp4_gpu4_node1_zero2_ds_offload")
      )
  
      suite.addTest(GPT2FuncTestCase("test_mp1_gpu1_node1_zero1"))
      suite.addTest(GPT2FuncTestCase("test_mp1_gpu2_node1_zero1"))
      suite.addTest(GPT2FuncTestCase("test_mp2_gpu4_node1_zero1"))
      suite.addTest(GPT2FuncTestCase("test_mp4_gpu4_node1_zero1"))
  
      suite.addTest(GPT2FuncTestCase("test_mp1_gpu1_node1_zero2"))
      suite.addTest(GPT2FuncTestCase("test_mp1_gpu2_node1_zero2"))
      suite.addTest(GPT2FuncTestCase("test_mp2_gpu4_node1_zero2"))
      suite.addTest(GPT2FuncTestCase("test_mp4_gpu4_node1_zero2"))
  
      suite.addTest(GPT2FuncTestCase("test_mp2_gpu4_node1_zero2_gas"))
  
      suite.addTest(GPT2FuncTestCase("test_optimizer_scheduler"))
  
      return suite
  
imports: import unittest
  import os
  import re
  from .test_common import BaseTestCase
label_desc: Use image node_id 2 for calling a global function with example usage: suite() and returns: suite