Add MNIST example with SPMD for JAX #2346
Conversation
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here.
Needs approval from an approver in each of these files:
Approvers can indicate their approval by writing |
google-oss-prow
bot
requested review from
andreyvelich,
tenzen-y and
terrytangyuan
There was a problem hiding this comment.
Thank you for this effort @sandipanpanda!
I left a few comments.
/cc @Electronic-Waste @kubeflow/wg-training-leads @kuizhiqing @mimowo @kannon92
| @@ -74,3 +74,7 @@ jobs: | |||
| platforms: linux/amd64 | |||
| dockerfile: examples/pytorch/deepspeed-demo/Dockerfile | |||
| context: examples/pytorch/deepspeed-demo | |||
| - component-name: jaxjob-mnist | |||
There was a problem hiding this comment.
Please can you also add the E2E tests for this example: https://github.com/kubeflow/training-operator/blob/master/sdk/python/test/e2e/test_e2e_jaxjob.py ?
| @@ -0,0 +1,27 @@ | |||
| FROM python:3.12 | |||
There was a problem hiding this comment.
Should we use alpine image here ?
There was a problem hiding this comment.
Made some adjustments to the current Dockerfile to use alpine image
FROM python:3.13-alpine
# Install build dependencies and other required tools using apk
RUN apk update && apk add --no-cache \
build-base \
cmake \
git \
glog-dev \
gflags-dev \
protobuf-dev \
linux-headers \
libc-dev \
&& rm -rf /var/cache/apk/*
# Upgrade pip and install jax and absl-py
RUN pip install --upgrade pip
RUN pip install --upgrade jax absl-py
# Clone and install gloo from source
RUN git clone https://github.com/facebookincubator/gloo.git \
&& cd gloo \
&& git checkout 43b7acbf372cdce14075f3526e39153b7e433b53 \
&& mkdir build \
&& cd build \
&& cmake ../ \
&& make \
&& make install
# Set working directory
WORKDIR /app
# Add application files
ADD datasets.py spmd_mnist_classifier_fromscratch.py /app
# Set entrypoint to run the script
ENTRYPOINT ["python3", "spmd_mnist_classifier_fromscratch.py"]
It increases build time by many times and throws error:
$ docker build -t sandipanify/jaxjob-mnist .
[+] Building 419.2s (9/11) docker:default
=> [internal] load build definition from Dockerfile 0.0s
=> => transferring dockerfile: 926B 0.0s
=> [internal] load metadata for docker.io/library/python:3.13-alpine 0.0s
=> [internal] load .dockerignore 0.7s
=> => transferring context: 2B 0.0s
=> [internal] load build context 0.1s
=> => transferring context: 89B 0.0s
=> CACHED [1/7] FROM docker.io/library/python:3.13-alpine 0.0s
=> [2/7] RUN apk update && apk add --no-cache build-base cmake git glog-dev gflags-dev protobuf-dev linux-headers libc-dev && rm -rf /var/cache/apk/* 42.0s
=> [3/7] RUN pip install --upgrade pip 4.6s
=> [4/7] RUN pip install --upgrade jax absl-py 347.0s
=> ERROR [5/7] RUN git clone https://github.com/facebookincubator/gloo.git && cd gloo && git checkout 43b7acbf372cdce14075f3526e39153b7e433b53 && mkdir build && cd build 24.4s
------
> [5/7] RUN git clone https://github.com/facebookincubator/gloo.git && cd gloo && git checkout 43b7acbf372cdce14075f3526e39153b7e433b53 && mkdir build && cd build && cmake ../ && make && make install:
0.570 Cloning into 'gloo'...
2.408 Note: switching to '43b7acbf372cdce14075f3526e39153b7e433b53'.
2.408
2.408 You are in 'detached HEAD' state. You can look around, make experimental
2.408 changes and commit them, and you can discard any commits you make in this
2.408 state without impacting any branches by switching back to a branch.
2.408
2.408 If you want to create a new branch to retain commits you create, you may
2.408 do so (now or later) by using -c with the switch command. Example:
2.408
2.408 git switch -c <new-branch-name>
2.408
2.408 Or undo this operation with:
2.408
2.408 git switch -
2.408
2.408 Turn off this advice by setting config variable advice.detachedHead to false
2.408
2.408 HEAD is now at 43b7acb Add removed tests to TARGETS
2.462 CMake Deprecation Warning at CMakeLists.txt:1 (cmake_minimum_required):
2.462 Compatibility with CMake < 3.10 will be removed from a future version of
2.462 CMake.
2.462
2.462 Update the VERSION argument <min> value. Or, use the <min>...<max> syntax
2.462 to tell CMake that the project requires at least <min> but has been updated
2.462 to work with policies introduced by <max> or earlier.
2.462
2.462
2.637 -- The CXX compiler identification is GNU 14.2.0
2.823 -- The C compiler identification is GNU 14.2.0
2.849 -- Detecting CXX compiler ABI info
2.972 -- Detecting CXX compiler ABI info - done
3.004 -- Check for working CXX compiler: /usr/bin/c++ - skipped
3.005 -- Detecting CXX compile features
3.007 -- Detecting CXX compile features - done
3.028 -- Detecting C compiler ABI info
3.145 -- Detecting C compiler ABI info - done
3.189 -- Check for working C compiler: /usr/bin/cc - skipped
3.190 -- Detecting C compile features
3.191 -- Detecting C compile features - done
3.201 -- Build type not set -- defaulting to Release
3.202 -- Gloo build as STATIC library
3.215 -- Configuring done (0.8s)
3.234 -- Generating done (0.0s)
3.235 -- Build files have been written to: /gloo/build
3.292 [ 2%] Building CXX object gloo/CMakeFiles/gloo.dir/algorithm.cc.o
4.226 [ 5%] Building CXX object gloo/CMakeFiles/gloo.dir/allgather.cc.o
5.447 [ 7%] Building CXX object gloo/CMakeFiles/gloo.dir/allgatherv.cc.o
6.993 [ 10%] Building CXX object gloo/CMakeFiles/gloo.dir/allreduce.cc.o
9.389 [ 12%] Building CXX object gloo/CMakeFiles/gloo.dir/allreduce_local.cc.o
10.82 [ 15%] Building CXX object gloo/CMakeFiles/gloo.dir/alltoall.cc.o
11.78 [ 17%] Building CXX object gloo/CMakeFiles/gloo.dir/alltoallv.cc.o
13.70 [ 20%] Building CXX object gloo/CMakeFiles/gloo.dir/barrier.cc.o
14.69 [ 22%] Building CXX object gloo/CMakeFiles/gloo.dir/broadcast.cc.o
15.86 [ 25%] Building CXX object gloo/CMakeFiles/gloo.dir/context.cc.o
17.20 [ 27%] Building CXX object gloo/CMakeFiles/gloo.dir/gather.cc.o
18.10 [ 30%] Building CXX object gloo/CMakeFiles/gloo.dir/gatherv.cc.o
19.36 [ 32%] Building CXX object gloo/CMakeFiles/gloo.dir/reduce.cc.o
20.72 [ 35%] Building CXX object gloo/CMakeFiles/gloo.dir/scatter.cc.o
21.78 [ 37%] Building CXX object gloo/CMakeFiles/gloo.dir/types.cc.o
22.35 [ 40%] Building CXX object gloo/CMakeFiles/gloo.dir/common/logging.cc.o
23.15 [ 42%] Building CXX object gloo/CMakeFiles/gloo.dir/common/utils.cc.o
23.47 [ 45%] Building CXX object gloo/CMakeFiles/gloo.dir/common/linux.cc.o
24.12 /gloo/gloo/common/linux.cc: In function 'int gloo::getInterfaceSpeedGLinkSettings(int, ifreq*)':
24.12 /gloo/gloo/common/linux.cc:201:20: error: '__caddr_t' was not declared in this scope; did you mean 'in_addr_t'?
24.12 201 | ifr->ifr_data = (__caddr_t)&ecmd;
24.12 | ^~~~~~~~~
24.12 | in_addr_t
24.12 /gloo/gloo/common/linux.cc: In function 'int gloo::getInterfaceSpeedGSet(int, ifreq*)':
24.12 /gloo/gloo/common/linux.cc:229:20: error: '__caddr_t' was not declared in this scope; did you mean 'in_addr_t'?
24.12 229 | ifr->ifr_data = (__caddr_t)&edata;
24.12 | ^~~~~~~~~
24.12 | in_addr_t
24.35 make[2]: *** [gloo/CMakeFiles/gloo.dir/build.make:317: gloo/CMakeFiles/gloo.dir/common/linux.cc.o] Error 1
24.35 make[1]: *** [CMakeFiles/Makefile2:216: gloo/CMakeFiles/gloo.dir/all] Error 2
24.35 make: *** [Makefile:136: all] Error 2
------
Dockerfile:20
--------------------
19 | # Clone and install gloo from source
20 | >>> RUN git clone https://github.com/facebookincubator/gloo.git \
21 | >>> && cd gloo \
22 | >>> && git checkout 43b7acbf372cdce14075f3526e39153b7e433b53 \
23 | >>> && mkdir build \
24 | >>> && cd build \
25 | >>> && cmake ../ \
26 | >>> && make \
27 | >>> && make install
28 |
--------------------
ERROR: failed to solve: process "/bin/sh -c git clone https://github.com/facebookincubator/gloo.git && cd gloo && git checkout 43b7acbf372cdce14075f3526e39153b7e433b53 && mkdir build && cd build && cmake ../ && make && make install" did not complete successfully: exit code: 2
Might be because Alpine Linux is a minimal distribution and some of the required libraries (like networking headers) might be missing or incompatible with the build. Issues similar to https://pythonspeed.com/articles/alpine-docker-python/
There was a problem hiding this comment.
We could use python:3.13-slim though as it is is based on Debian and is less likely to cause compatibility issues with C++ libraries like gloo
There was a problem hiding this comment.
Ok, maybe we can explore which C++ libs are needed to be installed in the future when using alpine image.
|
|
||
| WORKDIR /app | ||
|
|
||
| ADD datasets.py spmd_mnist_classifier_fromscratch.py /app |
There was a problem hiding this comment.
I would suggest to add ENTRYPOINT to the dockerfile, since it is a good practice to allow users to know the default command to run it. E.g. python3 spmd_mnist_classifier_fromscratch.py
| --- | ||
|
|
||
| ```bash | ||
| $ kubectl get pods -n kubeflow -l training.kubeflow.org/job-name=jaxjob-mnist |
There was a problem hiding this comment.
Please can you add the sample output of this command, e.g.:
$ kubectl get pods -n kubeflow -l training.kubeflow.org/job-name=jaxjob-mnist
NAME READY STATUS RESTARTS AGE
pod-name 3/3 Running 0 25d| ```bash | ||
| $ PODNAME=$(kubectl get pods -l training.kubeflow.org/job-name=jaxjob-mnist,training.kubeflow.org/replica-type=worker,training.kubeflow.org/replica-index=0 -o | ||
| name -n kubeflow) | ||
| $ kubectl logs -f ${PODNAME} -n kubeflow |
There was a problem hiding this comment.
The same here with the output
| @@ -0,0 +1,155 @@ | |||
| """An MNIST example with single-program multiple-data (SPMD) data parallelism. | |||
There was a problem hiding this comment.
Do we need to add licence here ?
| num_processes = int(os.getenv("NUM_PROCESSES")) | ||
| coordinator_address = os.getenv("COORDINATOR_ADDRESS") | ||
| coordinator_port = int(os.getenv("COORDINATOR_PORT")) | ||
| coordinator_address = f"{coordinator_address}:{coordinator_port}" |
There was a problem hiding this comment.
I think, you can directly populate the coordinator_address here, since you don't use port outside of the address.
| coordinator_address = f"{coordinator_address}:{coordinator_port}" | |
| coordinator_address = f"{os.getenv("COORDINATOR_ADDRESS")}:{int(os.getenv("COORDINATOR_PORT")}" |
|
@andreyvelich: GitHub didn't allow me to request PR reviews from the following users: mimowo, kannon92. Note that only kubeflow members and repo collaborators can review this PR, and authors cannot review their own PRs. In response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository. |
|
Hi @sandipanpanda, did you get a chance to finish this PR ? |
|
I am working on this. Will try to finish by today. |
Illustrate how to use JAX's `pmap` to express and execute single-program multiple-data (SPMD) programs for data parallelism along a batch dimension Signed-off-by: Sandipan Panda <samparksandipan@gmail.com>
Use -- server-side to install the latest local changes of Training Operator control plane Signed-off-by: Sandipan Panda <samparksandipan@gmail.com>
Signed-off-by: Sandipan Panda <samparksandipan@gmail.com>
sandipanpanda
force-pushed
the
jax-example
branch
from
8781809 to
e04bcf9
Compare
Pull Request Test Coverage Report for Build 12417960793Details
💛 - Coveralls |
|
It looks like E2E tests are failing, can you check it @sandipanpanda ? |
Signed-off-by: Sandipan Panda <samparksandipan@gmail.com>
|
@andreyvelich It seems to be failing because |
It would be nice if we could build an image via CI and load into the Kind cluster, so we can build image directly from the code that is submitted in the PR. |
Illustrate how to use JAX's
pmapto express and execute single-program multiple-data (SPMD) programs for data parallelism along a batch dimensionAdapted from: https://github.com/jax-ml/jax/blob/main/examples/spmd_mnist_classifier_fromscratch.py
Output: https://gist.github.com/sandipanpanda/1c92e9875e7628338d39251f5e1aa673