Function main

tests/sharding_dump.py:258–288

Load a config that describes a model with topology and slices to be dumped.

(argv: Sequence[str])

def main(argv: Sequence[str]) -> None:
  """Load a config that describes a model with topology and slices to be dumped."""
  jax.config.update("jax_default_prng_impl", "unsafe_rbg")
  os.environ["LIBTPU_INIT_ARGS"] = (
      os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
  )
  print("Starting sharding_tests.py...", flush=True)

  config = pyconfig.initialize(argv)
  validate_config(config)

  json_path = (
      f"sharding_info/{config.model_name}/"
      f"{config.compile_topology}/"
      f"slice_{config.compile_topology_num_slices}/"
      f"named_shardings.json"
  )

  try:
    topology_mesh = get_topology_mesh(config)
    _, _, state_mesh_shardings, _ = get_shaped_inputs(topology_mesh, config)
  except:  # pylint: disable=bare-except
    state_mesh_shardings = {}

  if state_mesh_shardings == {}:
    return

  sharding_dict = named_shardings_to_json(state_mesh_shardings)
  save_named_sharding_dict(json_path, sharding_dict)
  load_named_sharding_json(json_path)
  print(config.model_name, config.compile_topology)


if __name__ == "__main__":