PAX Check loss and perf

[maanug-nv]

• Adds baselines for 4 of the PAX mgmn tests
• Adds pytest files to compare tensorboard file result of mgmn tests against baselines
• Adds job to run above pytest in workflow

[terrykong]

It would be neat to surface the results of this either in a GITHUB_STEP_SUMMARY or a badge somewhere. Any thoughts about adding that?

[maanug-nv]

Agreed, that's my next step 😄. Just wanted to split into two PRs.

[yhtang]

How are the baseline numbers in the JSON files derived? We need to document the source of the data so that when a significant change happens, we know where to look for the diffs.

If they are from a particular Workflow run, maybe we can simply reference it and checkout the associated artifacts? This way, the source and methodology of the data is self-explanatory.

[terrykong]

@yhtang Raises a good point about baseline provenance.

What about in addition to the python file that create the baseline jsons from artifacts, you add a script that takes jobids and pulls the artifacts and calls this python script? That way from the jobid, we can encode the URL of the workflow run and the date it ran in the JSON for inspection and repeatability?

[yhtang]

What about in addition to the python file that create the baseline jsons from artifacts, you add a script that takes jobids and pulls the artifacts and calls this python script? That way from the jobid, we can encode the URL of the workflow run and the date it ran in the JSON for inspection and repeatability?

I like this idea. Essentially, we can wrap this up as a reusable workflow, say _compare_metrics.yaml, which we can mentally regard as a function. This workflow/function will accept 2 workflow/job IDs, download the associated artifacts, and determine if the differences between certain metrics (can be specified via an input) is large.

This way, we can reuse this workflow for comparing results for PAX, T5X, JAX, TE, etc...