The manifest rebase #406
The manifest rebase #406
Conversation
author Yu-Hang Tang <Tang.Maxin@gmail.com> 1698050497 +0000 committer Terry Kong <terryk@nvidia.com> 1701417045 -0800 pip-compile changes Updated t5-large perf (#342) Update Pax README and sub file (#345) - Adds FP8 documentation - Updates perf table - Makes some other minor improvements for readability Adds CUDA_MODULE_LOADING=EAGER to core jax container env vars (#329) Re-enable NVLS in nightly containers (#331) NVLS was disabled due to a known issue in NCCL 2.17 that caused intermittent hangs. The issue has been resolved in NCCL 2.18, so we are safe to re-enable NVLS. --------- Co-authored-by: Terry Kong <terryk@nvidia.com> Update Pax TE patch to point to rebased branch (#348) Loosens t5x loss tests relative tolerances (#343) Relaxing the relative tolerance on the loss tests since it was leading to too many false positives. For reference, deviation in loss for the t5 model can sometimes be up to 15% at the start of training with real data. Adds rosetta-t5x TE + no-TE tests that enable the correct configs for testing (#332) - [ ] Add capability to retroactively test with newer test-t5x.sh like in [t5x-wget-test](https://github.com/NVIDIA/JAX-Toolbox/tree/t5x-wget-test) - [ ] Sets `ENABLE_TE=1` in the Dockerfile.t5x which is identical to the logic from before where it was always enabled in rosetta-t5x Fix markdown hyperlink for jax package on frontpage readme (#319) Adds a --seed option to test-t5x.sh to ensure determinism (#344) To ensure that the tests results for a particular container are reproducible between runs, this change introduces a seed argument that sets the jax seed and dataset seed to 42. It remains configurable, but now there shouldn't be variance given the same container. - Also fixes a typo where --steps-per-epoch wasn't in the usage doc of this script Co-authored-by: NVIDIA <jax@nvidia.com> Co-authored-by: Yu-Hang "Maxin" Tang <Tang.Maxin@gmail.com> Dynamic workflow run names (#356) This change introduces the dynamic [run name field](https://github.blog/changelog/2022-09-26-github-actions-dynamic-names-for-workflow-runs/#:~:text=GitHub%20Actions%20customers%20can%20now,visit%20the%20GitHub%20Actions%20community.) `run-name`. It's currently difficult on mobile to find the "workflow_run" that corresponds to a particular date, so hopefully this helps identify which builds were nightly vs which builds were manually triggered. I couldn't find a good way to dynamically look up the `name` field, so for now I copied all of names. I also wasn't able to find a "created_at" for the scheduled workflows, so those don't have timestamps for now. __Assumptions__: * "workflow_run" == nightly since "scheduled" events only happen on `main` and `workflow_run` are only run for concrete workflows and not reusable workflows - [x] Test the workflow_run codepath - [x] Test the scheduled codepath  Fix random failling tests for backend_independent on V100 (#351) Fixes randomly failures in the backend-independent section of JAX unit tests: ``` Cannot find a free accelerator to run the test on, exiting with failure ``` Changes: limit the number of concurrent test jobs even for backend-independent tests, which do create GPU contexts. As a clarification, `--jobs` and `--local_test_jobs` do not make a difference for our particular CI pipeline, since JAX is built in a separate CI job anyway. References (From Reed Wanderman-Milne @ Google): > 1. In particular, you have to set NB_GPUS, JOBS_PER_ACC, and J correctly or you can get that error (I recently got the same error by not setting those correctly) > 2. (also I think --jobs should be --local_test_jobs in that code block, no reason to restrict the number of jobs compiling JAX) Propagate error code in ViT tests (#357) Merges rosetta unit tests and takes off the marker which spun up another matrix job (#360) This should simplify the rosetta tests and save some time since another matrix job was started for one test Propagate build failures (#363) Always run the `publish-build` step, regardless of whether the rosetta pax/t5x build was attempted. This ensures that badges correctly reflect build failures due to dependent builds failing. Patch for JAX core container (ARM64) (#367) Add patch to XLA to be able to build JAX core container for ARM64 Update the doc for USE_FP8 (#349) This PR provides guidance on how to use the new configuration option, `USE_FP8`, to enable native FP8 support on Hopper GPUs. Update the native-fp8 guide with cudnn layer norm (#368) This PR updates the guide to include the new flag to enable the cudnn layer norm. cc. @ashors1 @terrykong @nouiz Add WAR for XLA NCCL bug causing OOMs (#362) A stopgap for #346 fix TE multi-device test fix lzma build issue edit TE test name fix TE arm64 test install error remove --install option from get-source.sh fix TE arm64 test install error disable sandbox i'm jet-lagged use Pax image for TE testing Fix job dependency
generated patch dir comment Add steps to archive patches in run Date the patches for readability Better log msg switch to --3way since that produces a merge conflict to help understand the conflict Switch to mealkit+finalize mechanic for rosetta builds Add github.run_id to artifacts for provenance Update all rosetta workflows with mealkit/final mechanism
author Terry Kong <terryk@nvidia.com> 1700265014 -0800 committer Terry Kong <terryk@nvidia.com> 1701417338 -0800 parent 8a43f4a author Terry Kong <terryk@nvidia.com> 1700265014 -0800 committer Terry Kong <terryk@nvidia.com> 1701417298 -0800 Adds (1) bump.sh which bumps the manifest and pins the patches (2) updates create-distribution.sh to work with manifests (3) move everything to .github/container sandbox fix write add propagation of trial branch to all workflows and update sandbox to test synchronous workflow check wip test wip changes wip wip don't need wip wip remove make trial branch contingent on publishing Update get-source and initial update for jax build to accept manifest update manifest jax build partially working + patches update pax/t5x dockerfiles, add more repos into manifest, and update pip-finalize to use *.in instead of manifest.txt update manifest with rest of repos and patches missing arg fix jax/pax/t5x all builds work now! update manifest file everywhere fix all workflows cleanup get the context right fix all broken tests custom pip distribution works
| git add -u | ||
| git commit -m 'Adds JAX_TOOLBOX_VCS_EQUIVALENCY as a trigger to treat all github VCS installs for a package as equivalent. The spec of the last encountered version will be used' | ||
| EOF | ||
| RUN pip install --upgrade --no-cache-dir -e /opt/pip pip-tools && rm -rf ~/.cache/* |
There was a problem hiding this comment.
Would it cause an issue for the end user if we use a non-standard version of pip? For example, some users may want to treat VCS and PyPI/local packages as different. Perhaps we can either:
- turn the VCS equivalence into an option that defaults to false
- re-install upstream pip as the last step of container build
?
There was a problem hiding this comment.
Right now this feature is behind an environment variable JAX_TOOLBOX_VCS_EQUIVALENCY=true and that's only passed in pip-finalize.sh. If left un-set, it's equivalent to upstream. Do you think that is sufficient?
There was a problem hiding this comment.
Ah I see. That should be sufficient. BTW, another point to note is that given that the baking (pip-compile + pip-sync) step will be independent between the upstream and rosetta containers, it is technically possible to let the upstream JAX/T5X/PAX/etc packages use free-floating VCS installs and rosetta containers use pinned VCS installs. Whether it will bring in any advantage is another question, though.
| ARG REF_XLA=main | ||
| ARG REF_FLAX=main | ||
| ARG REF_TE=main | ||
| ARG SRC_MANIFEST_FILE=manifest.yaml |
There was a problem hiding this comment.
I would prefer we place the manifest file parsing mechanism in the CI code, while maintaining simple Dockerfiles. Otherwise the Dockerfile will no longer be self-contained and less reusable by external users.
There was a problem hiding this comment.
Don't we want to be prescriptive here since the jax ecosystem is difficult to pin down and we don't want external users to run into the same issues as us? In this case the prescription is: an external user has to use a manifest file to build the jax-toolbox containers.
I agree we want the dockerfiles to be simple and self-contained. I guess "simple" here is subjective, but I think the change to only require a two build args + COPY to get the manifest in, simplifies and takes away the many --build-args one would need to build any of the images. Could you elaborate on what I can make simpler?
| RUN --mount=type=ssh \ | ||
| --mount=type=secret,id=SSH_KNOWN_HOSTS,target=/root/.ssh/known_hosts \ | ||
| git clone "${REPO_XLA}" "${SRC_PATH_XLA}" && cd "${SRC_PATH_XLA}" && git checkout ${REF_XLA} | ||
| get-source.sh -l xla -m ${MANIFEST_FILE} |
There was a problem hiding this comment.
Similarly, get-source should be preferrably self-contained.
There was a problem hiding this comment.
Are you asking for this to be configurable with build-args?
| @@ -0,0 +1,97 @@ | |||
| #!/bin/bash | |||
There was a problem hiding this comment.
bump.sh -> .github/workflows/scripts/bump-manifest.sh so that the logic remains in CI?
| @@ -0,0 +1,103 @@ | |||
| # Updated in: 2023-12-01 | |||
There was a problem hiding this comment.
Move to either .github/workflows or just root?
No description provided.