add python profiler

[sfc-gh-sbekman]

This PR will allow you to add to the launcher --python_profile <sortkey>, which will run the code under cProfile, where sort keys are tottime or cumtime and then at the end of the run it will dump to stdout something like:

         2864954 function calls (2783834 primitive calls) in 6.499 seconds

   Ordered by: cumulative time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
   515/18    0.002    0.000    3.857    0.214 parameter_offload.py:300(_post_forward_module_hook)
      290    0.001    0.000    2.951    0.010 autocast_mode.py:512(decorate_fwd)
       10    0.000    0.000    2.863    0.286 loss_utils.py:35(LigerForCausalLMLoss)
       10    0.000    0.000    2.862    0.286 loss_utils.py:9(fixed_fused_linear_cross_entropy)
       10    0.000    0.000    2.862    0.286 functional.py:55(liger_fused_linear_cross_entropy)
       10    0.000    0.000    2.667    0.267 fused_linear_cross_entropy.py:208(forward)
       10    1.122    0.112    2.667    0.267 fused_linear_cross_entropy.py:16(fused_linear_cross_entropy_forward)
       10    0.000    0.000    2.210    0.221 debug.py:133(get_mem_metrics)

I'm marking this API as experimental, since it's possible we might want to get the profile data into a file and then add per rank profile and more. So this is a very rough solution at the moment which I think will evolve.

Note we can't use yaml config here since it's too early in the game, hence the argparse API.

[sfc-gh-truwase]

I don't know how dependencies are managed in AT, but does cProfile need to be added as dependency?

[sfc-gh-sbekman]

it's distributed with python core https://docs.python.org/3/library/profile.html

dependencies in AT are managed:

1. required via pyproject.toml
2. optional via code requiring it doing the checking - like liger-kernel