Function main

src/MaxText/train_compile.py:197–265

(argv: Sequence[str])

def main(argv: Sequence[str]) -> None:
  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 train_compile.py...", flush=True)

  # Parse and validate configuration
  config = pyconfig.initialize(argv)
  validate_config(config)

  # Create target mesh
  topology_mesh = get_topology_mesh(config)

  # Print system information after building the compile topology to avoid
  # prematurely initializing the backend.
  max_utils.print_system_information()

  # Get shaped inputs
  shaped_train_args, shaped_train_kwargs, state_mesh_shardings, model = get_shaped_inputs(topology_mesh, config)

  # Get data sharding
  data_sharding = sharding.get_input_data_sharding(config, topology_mesh)

  # Get function to compile and shardings
  func_to_compile, in_shard, out_shard, static_argnums, donate_argnums = (
      maxtext_utils.get_functional_train_with_signature(
          train.train_step, data_sharding, state_mesh_shardings, model, config
      )
  )

  # print weights sharding info under debug sharding mode
  if config.debug_sharding:
    max_utils.print_non_trivial_mesh_axis(topology_mesh)
    maxtext_utils.print_state_mesh_shardings_params(shaped_train_args[0], state_mesh_shardings, topology_mesh)

  # Compile
  print("Jitting and compiling train step...", flush=True)
  compiled = jit_and_compile(
      func_to_compile,
      shaped_train_args,
      shaped_train_kwargs,
      topology_mesh,
      in_shard,
      out_shard,
      static_argnums,
      donate_argnums,
      nn_partitioning.axis_rules(config.logical_axis_rules),
  )
  print("Jitting and compilation complete!", flush=True)

  # Serialize and save the compiled object
  if config.compiled_trainstep_file != "":
    print("Saving compiled object...")
    save_compiled(compiled, config.compiled_trainstep_file)
    print(f"Successfully saved compiled object as {config.compiled_trainstep_file}")
  print("Finished train_compile.py successfully!", flush=True)
  print(f"Cost analysis: {compiled.cost_analysis()}")