|
| 1 | +# MaxDiffusion on DGX Spark GPU: A complete User Guide |
| 2 | + |
| 3 | +This guide provides a detailed step-by-step walkthrough for setting up and running the maxdiffusion library within a custom Docker environment on an ARM-based machine with NVIDIA GPU support (like a spark VM). We will cover everything from building the optimized Docker image to generating your first image and retrieving it successfully. |
| 4 | + |
| 5 | +## Prerequisites |
| 6 | + |
| 7 | +Before you begin, ensure you have the following: |
| 8 | + |
| 9 | +- Access to an ARM-based Linux machine (e.g., your spark VM) with Docker installed. |
| 10 | +- NVIDIA GPU drivers and the NVIDIA Container Toolkit installed on the host machine to enable `--gpus all` support. |
| 11 | +- The maxdiffusion source code cloned onto the machine. |
| 12 | + - Branch: dgx_spark |
| 13 | +- A requirements.txt file in the root of the maxdiffusion directory |
| 14 | +- An internet connection for the initial Docker build and for downloading models (if not cached). |
| 15 | + |
| 16 | +## Part 1: Building the Optimized Docker Image |
| 17 | + |
| 18 | +The foundation of a smooth workflow is a well-built Docker image. The following Dockerfile is optimized for build speed by caching dependencies, ensuring that code changes don't require a full reinstall of all libraries. |
| 19 | + |
| 20 | +### Step1: Create the Dockerfile |
| 21 | + |
| 22 | +In the root directory of your maxdiffusion project, create a file named box.Dockerfile and paste the following content into it. |
| 23 | + |
| 24 | +```docker |
| 25 | +# Base image for ARM64 with CUDA support |
| 26 | +FROM nvcr.io/nvidia/cuda-dl-base@sha256:3631d968c12ef22b1dfe604de63dbc71a55f3ffcc23a085677a6d539d98884a4 |
| 27 | +
|
| 28 | +# Set environment variables (these rarely change) |
| 29 | +ENV PIP_BREAK_SYSTEM_PACKAGES=1 |
| 30 | +ENV DEBIAN_FRONTEND=noninteractive |
| 31 | +
|
| 32 | +# Install system-level dependencies (these change very infrequently) |
| 33 | +RUN apt-get update && apt-get install -y python3 python3-pip |
| 34 | +RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 1 && \ |
| 35 | + update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1 |
| 36 | +
|
| 37 | +WORKDIR /app |
| 38 | +
|
| 39 | +# --- Dependency Installation Layer --- |
| 40 | +# First, copy only the requirements file to leverage caching |
| 41 | +COPY requirements.txt . |
| 42 | +
|
| 43 | +# Install dependencies from requirements.txt |
| 44 | +RUN pip install -r requirements.txt |
| 45 | +
|
| 46 | +# Install other major Python libraries in separate layers for better caching |
| 47 | +RUN pip install --upgrade torch torchvision |
| 48 | +RUN pip install "jax[cuda13-local]==0.7.2" |
| 49 | +
|
| 50 | +# Install Transformer Engine and its dependency pybind11 |
| 51 | +RUN pip install pybind11 && \ |
| 52 | + export NVTE_FRAMEWORK=jax && \ |
| 53 | + pip install --no-build-isolation 'transformer_engine[jax]==2.6.0' |
| 54 | +
|
| 55 | +# --- Application Code Layer --- |
| 56 | +# Now, copy your application source code. This layer is rebuilt only when your code changes. |
| 57 | +COPY . . |
| 58 | +
|
| 59 | +# Install the maxdiffusion package from the copied source |
| 60 | +RUN pip install . |
| 61 | +
|
| 62 | +# Set a default command to keep the container running for interactive use |
| 63 | +CMD ["/bin/bash"] |
| 64 | +``` |
| 65 | + |
| 66 | +### Step2: Build the Image |
| 67 | + |
| 68 | +Open your terminal on the spark VM, navigate to the root directory of the maxdiffusion project, and run the build command: |
| 69 | + |
| 70 | +```bash |
| 71 | +docker build -f box.Dockerfile -t maxdiffusion-arm-gpu . |
| 72 | +``` |
| 73 | + |
| 74 | +This command will execute the steps in your Dockerfile, download the necessary layers, install all dependencies, and create a local Docker image named maxdiffusion-arm-gpu. The first build may take some time. Subsequent builds will be much faster if you only change the source code. |
| 75 | + |
| 76 | +## Part 2: Running the Container for Image Generation |
| 77 | + |
| 78 | +To run the image generator effectively, we need to connect our local machine's folders to the container. This prevents re-downloading models and makes it easy to retrieve the output images. |
| 79 | + |
| 80 | +### Step 1: Create a Local Output Directory |
| 81 | + |
| 82 | +On your spark VM, create a directory to store the generated images. |
| 83 | + |
| 84 | +```bash |
| 85 | +mkdir -p ~/maxdiffusion_output |
| 86 | +``` |
| 87 | + |
| 88 | +### Step 2: Launch the Container with Volume Mounts |
| 89 | + |
| 90 | +Run the following command to start an interactive session inside your container. This command links your Hugging Face cache (to avoid re-downloading models) and the output directory you just created. |
| 91 | + |
| 92 | +```bash |
| 93 | +docker run -it --gpus all \ |
| 94 | +-v ~/.cache/huggingface:/root/.cache/huggingface \ |
| 95 | +-v ~/maxdiffusion_output:/tmp \ |
| 96 | +maxdiffusion-arm-gpu |
| 97 | +``` |
| 98 | +Your terminal prompt will change, indicating you are now inside the running container. |
| 99 | + |
| 100 | +## Part 3: Generating Your First Image |
| 101 | + |
| 102 | +Now that you are inside the container's interactive shell, you can execute the image generation script. Run the following command: |
| 103 | + |
| 104 | +```bash |
| 105 | +NVTE_FRAMEWORK=JAX NVTE_FUSED_ATTN=1 HF_HUB_ENABLE_HF_TRANSFER=1 python src/maxdiffusion/generate_flux.py src/maxdiffusion/configs/base_flux_dev.yml jax_cache_dir=/tmp/cache_dir run_name=flux_test output_dir=/tmp/ prompt='A cute corgi lives in a house made out of sushi, anime' num_inference_steps=28 split_head_dim=True per_device_batch_size=1 attention="cudnn_flash_te" hardware=gpu |
| 106 | +``` |
| 107 | +The script will initialize, use the models from your mounted cache, and begin the generation process. |
| 108 | + |
| 109 | +## Part 4: Accessing Your Generated Image |
| 110 | + |
| 111 | +The generation script saves the final image to its working directory (/app) inside the container. Here is the complete workflow to get that image onto your Mac. |
| 112 | + |
| 113 | +### Step 1: Copy the Image from Container to VM |
| 114 | + |
| 115 | +Open a new terminal window and connect to your spark VM. Do not close the terminal where the container is running. |
| 116 | +First, find your container's ID: |
| 117 | + |
| 118 | +```bash |
| 119 | +docker ps |
| 120 | +``` |
| 121 | + |
| 122 | +Look for the container with the image maxdiffusion-arm-gpu and note its ID (e.g., 9049895399fc). |
| 123 | +Now, copy the image from the container to a temporary location on the VM and fix its permissions. |
| 124 | + |
| 125 | +```bash |
| 126 | +# Copy the file to the /tmp/ directory on the VM |
| 127 | +docker cp 9049895399fc:/app/flux_0.png /tmp/flux_0.png |
| 128 | + |
| 129 | +# Change the file's owner to your user to avoid permission errors |
| 130 | +sudo chown parambole:parambole /tmp/flux_0.png |
| 131 | +``` |
| 132 | + |
| 133 | +### Step 2: Copy the Image from VM to Your MAC |
| 134 | + |
| 135 | +Now, open the Terminal app on your Mac and use the scp (secure copy) command to download the file from the VM. |
| 136 | + |
| 137 | +```bash |
| 138 | +scp parambole@spark:/tmp/flux_0.png . |
| 139 | +``` |
| 140 | + |
| 141 | +This command will download flux_0.png to the current directory on your Mac. You can now view your generated image! |
| 142 | + |
| 143 | +## Troubleshooting and Common Pitfalls |
| 144 | + |
| 145 | +Here are solutions to common issues you might encounter: |
| 146 | +- Error: `pip: command not found` during Docker build. |
| 147 | + - **Cause**: The base Docker image doesn't have pip in the system's default PATH. |
| 148 | + - **Solution**: The provided Dockerfile fixes this by explicitly installing python3-pip and using update-alternatives to create the necessary symbolic links. |
| 149 | +- Error: `externally-managed-environment` during `pip install`. |
| 150 | + - **Cause**: Newer versions of Debian/Ubuntu protect system Python packages from being modified by pip. |
| 151 | + - **Solution**: The `ENV PIP_BREAK_SYSTEM_PACKAGES=1` line in the `Dockerfile` safely bypasses this protection within the container's isolated environment. |
| 152 | +- Error: `OSError: ...is not a local folder and is not a valid model identifier` |
| 153 | + - **Cause**: The script is trying to download models from the Hugging Face Hub because it cannot find them locally. |
| 154 | + - **Solution**: This is solved by launching the container with the `-v ~/.cache/huggingface:/root/.cache/huggingface` flag, which gives the container access to your local model cache. |
| 155 | +- Error: `open ... permission denied` when trying to access a copied file. |
| 156 | + - **Cause**: Files copied from a Docker container with docker cp are owned by the root user by default. |
| 157 | + - **Solution**: After copying the file to the VM, immediately run `sudo chown your_user:your_user /path/to/file` to take ownership before trying to access or transfer it. |
| 158 | +- Can't find the generated image. |
| 159 | + - **Cause**: The script may not be saving the image to the directory specified by the output_dir argument. |
| 160 | + - **Solution**: Always check the script's source code to confirm the final save location. As we discovered, generate_flux.py saves to the current working directory (/app), not /tmp. Knowing this allows you to copy the file from the correct location. |
| 161 | +- If a process requires more memory than the available RAM, your system will crash with an "Out-of-Memory" (OOM) error. |
| 162 | + - `Swap memory is your safety net.` It's a designated space on your hard drive that the operating system uses as a "virtual" extension of your RAM. When RAM is full, the system moves less active data to the slower swap space, freeing up RAM for the immediate task. While it's slower than RAM, it's infinitely better than a system crash, ensuring your long-running training or generation jobs can complete successfully. For a machine with 119GB of RAM, adding 64GB of swap provides a robust buffer for memory-intensive operations. |
| 163 | + - Step 1: Create a 64GB Swap File |
| 164 | + - Run these commands on your spark-1c91 VM to create, format, and enable a permanent 64GB swap file. |
| 165 | + |
| 166 | + ```bash |
| 167 | + # Instantly allocate a 64GB file |
| 168 | + sudo fallocate -l 64G /swapfile |
| 169 | + # Set secure permissions (only root can access) |
| 170 | + sudo chmod 600 /swapfile |
| 171 | + # Format the file as swap space |
| 172 | + sudo mkswap /swapfile |
| 173 | + # Enable the swap file for the current session |
| 174 | + sudo swapon /swapfile |
| 175 | + # Add the swap file to the system's startup configuration to make it permanent |
| 176 | + echo '/swapfile none swap sw 0 0' | sudo tee -a /etc/fstab |
| 177 | + ``` |
| 178 | + |
| 179 | + - Step 2: Verify Swap is Active |
| 180 | + - Check that the swap space is correctly configured. |
| 181 | + |
| 182 | + ```bash |
| 183 | + free -h |
| 184 | + # The output should now show a 64GB total for Swap. |
| 185 | + ``` |
0 commit comments