Add profiling option to the fuji running script (#1491)

AXLearn has a jax profiling option, that works by giving
`--jax_profiler_port=9999` to the submission/deployment script.
This PR adds the option `--trace_steps`, a list of steps that we can
give to `fuji-train-perf.py` to get profiles at different steps.

Author: Steboss

.github/container/fuji-train-perf.py

@@ -127,6 +127,13 @@
     help="Final log file.",
 )
 
+parser.add_argument(
+    "--trace_steps",
+    type=list,
+    default=None,
+    help="Steps to trace (e.g. [1, 20, 50]). To profile the training give `--jax_profiler_port 9999` to the script.",
+)
+
 parser.add_argument("--world_size", type=int, help="Total number of GPUs")
 
 
@@ -245,6 +252,7 @@ def main(parsed_args):
     save_checkpoint_steps = parsed_args.save_checkpoint_steps
     write_summary_steps = parsed_args.write_summary_steps
     output_log_file = parsed_args.output_log_file
+    trace_steps = parsed_args.trace_steps
     world_size = parsed_args.world_size
 
     print(
@@ -306,6 +314,10 @@ def main(parsed_args):
     trainer_config.checkpointer.save_policy.n = save_checkpoint_steps
     trainer_config.checkpointer.keep_every_n_steps = save_checkpoint_steps
     trainer_config.summary_writer.write_every_n_steps = write_summary_steps
+
+    if trace_steps is not None:
+        trainer_config.start_trace_steps = trace_steps
+
     # Call AXLearn Jax setup
     launch.setup()
     # Setup the config