Skip to content

test tree_util changes #1396

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 23 commits into from
May 8, 2025
Merged

test tree_util changes #1396

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
620682a
test tree_util changes
Steboss Apr 16, 2025
8ee3620
Add XLA flag to avoid error
Steboss Apr 17, 2025
8d5b54a
@olupton comments
Steboss Apr 17, 2025
69ae1ac
fix manifest
Steboss Apr 17, 2025
49dffe2
fix manifest2
Steboss Apr 17, 2025
af902af
Test with xla commit 520f
Steboss Apr 24, 2025
40d92f4
fix buffer
Steboss Apr 24, 2025
aee4493
try newest version of xla
Steboss Apr 28, 2025
2c7f1d6
stupid error on git
Steboss Apr 28, 2025
95a250d
Merge branch 'main' into sbosisio/axlearn-test-tree
Steboss Apr 28, 2025
1812768
pin a specific jax commit
Steboss Apr 28, 2025
3b5a2e7
Merge branch 'sbosisio/axlearn-test-tree' of github.com:NVIDIA/JAX-To…
Steboss Apr 28, 2025
dd96fa1
pin jax version
Steboss Apr 28, 2025
67f831f
test axlearn with new modifications
Steboss Apr 28, 2025
c9ae84b
rebuild jax, test the error in utility, check xla status as well
Steboss May 6, 2025
1d166f8
Merge branch 'main' into sbosisio/axlearn-test-tree
Steboss May 6, 2025
026d06c
reset xla to original
Steboss May 6, 2025
1161aa1
Merge branch 'sbosisio/axlearn-test-tree' of github.com:NVIDIA/JAX-To…
Steboss May 6, 2025
c37df36
test a working branch
Steboss May 6, 2025
2b8a931
reset ci
Steboss May 6, 2025
90c89e4
fix ci
Steboss May 6, 2025
d8ab194
Update .github/container/manifest.yaml
Steboss May 7, 2025
1717757
Merge branch 'main' into sbosisio/axlearn-test-tree
Steboss May 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/container/Dockerfile.axlearn
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# syntax=docker/dockerfile:1-labs
ARG BASE_IMAGE=ghcr.io/nvidia/jax-mealkit:jax
ARG URLREF_AXLEARN=https://github.com/Steboss/axlearn.git#main
ARG URLREF_AXLEARN=https://github.com/apple/axlearn.git
ARG SRC_PATH_AXLEARN=/opt/axlearn

###############################################################################
Expand Down
5 changes: 5 additions & 0 deletions .github/container/manifest.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,3 +136,8 @@ pathwaysutils:
tracking_ref: main
latest_verified_commit: 359776d454940ffaa337c36d1df16308d44a95a9
mode: pip-vcs
axlearn:
url: https://github.com/Steboss/axlearn.git
tracking_ref: sbosisio/tree_util
latest_verified_commit:
mode: git-clone
2 changes: 2 additions & 0 deletions .github/eks-workflow-files/axlearn/axlearn-fuji-model.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@ spec:
mkdir -p ${TRAINER_DIR}
OUTPUT_LOG_FILE=${TRAINER_DIR}/output.log

export XLA_FLAGS="--xla_gpu_enable_command_buffer="

python3 /usr/local/bin/fuji-train-perf.py \
--module=text.gpt.c4_trainer \
--config=${CONFIG} \
Expand Down
4 changes: 3 additions & 1 deletion .github/workflows/_ci.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,8 @@ jobs:
CONTAINER_NAME: axlearn
DOCKERFILE: .github/container/Dockerfile.axlearn
RUNNER_SIZE: large
EXTRA_BUILD_ARGS: |
URLREF_AXLEARN=${{ fromJson(inputs.SOURCE_URLREFS).AXLEARN }}
secrets: inherit

collect-docker-tags:
Expand Down Expand Up @@ -521,7 +523,7 @@ jobs:
# merge the log files
cat \
/log/pytest-report-L0-unittest.jsonl
/log/pytest-report-L0-distributed-unittest.jsonl
/log/pytest-report-L0-distributed-unittest.jsonl
> /log/pytest-report.jsonl
EOF
STATISTICS_SCRIPT: |
Expand Down
Loading