code: def run_test_add_param_group(rank, world_size, tempfile_name):
      dist_init(rank, world_size, tempfile_name)
  
      # Test with all parameters trainable to begin with
      def all_trainable():
          params = []
          sizes = [9, 7, 5, 3]
          sizes_world = sizes * world_size
          for size in sizes_world[:-1]:
              params.append(torch.rand(size, 1))
  
          # Make sure that the params are trainable, enforces size-based partitioning
          for p in params:
              p.requires_grad = True
  
          o = optim.OSS(params, lr=0.1)
  
          assert len(o.param_groups) == 1
          o.add_param_group({"params": [torch.rand(3, 1)]})
  
          assert len(o.param_groups) == 2
  
          # Verify that added group is added to the correct partition making all have the same number of elements
          assert sum(
              [
                  x.numel()
                  for g in o.optim.param_groups
                  for x in g["params"]
              ]
          ) == sum(sizes)
          assert len(o.optim.param_groups) == 2
  
      # Test a pathological config with a first big non-trainable param
      def some_trainable():
          params = []
          for size in [100, 3, 5, 2, 6, 4]:
              params.append(torch.rand(size, 1))
  
          # Make sure that the params are trainable, enforces size-based partitioning
          for p in params[1:]:
              p.requires_grad = True
  
          o = optim.OSS(params, lr=0.1)
  
          assert len(o.param_groups) == 1
          o.add_param_group({"params": [torch.rand(3, 1)]})
  
          assert len(o.param_groups) == 2
          assert len(o.optim.param_groups) == 2
  
      all_trainable()
      some_trainable()
  
      dist.destroy_process_group()
  
imports: import copy
  from math import inf
  import tempfile
  from typing import Any, Dict, Type, cast
  import unittest
  import numpy
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  from torch.nn.parallel import DistributedDataParallel
  from fairscale.fair_dev.testing.testing import check_same_model_params, check_same_models_across_ranks, skip_if_no_cuda, skip_if_py39_no_cuda, skip_if_single_gpu
  from fairscale.internal import torch_version
  import fairscale.optim
label_desc: Use image node_id 3 for calling a global function with example usage: run_test_add_param_group(rank, world_size, tempfile_name) without return types