Add containerized DGX Spark / ARM64+Blackwell deployment and portability fixes#10
Open
finn1901 wants to merge 3 commits into
Open
Add containerized DGX Spark / ARM64+Blackwell deployment and portability fixes#10finn1901 wants to merge 3 commits into
finn1901 wants to merge 3 commits into
Conversation
…taset tooling Enables StyleDrive to build and run on NVIDIA DGX Spark (Grace-Blackwell, sm_121, ARM64), which the pinned x86_64/CUDA-11 stack cannot support. - Dockerfile + .dockerignore: NGC PyTorch ARM64 base (torch 2.7 / CUDA 12.x, Blackwell kernels); relaxes pins that collide with the base image's locked stack (numpy, setuptools, scikit-learn, protobuf, opencv-python, guppy3) and bumps hydra-core to >=1.3.2 (1.2.0 cannot import on Python >=3.11). Data is kept out of the image and bind-mounted at /data. - download/prepare_dataset.sh: resumable download + assembly into the expected dataset layout; fixes the meta_datas/sensor_blobs nesting, the HuggingFace /blob/ -> /resolve/ style-JSON URL, and validates the JSON actually parsed. - worker/ray_distributed_no_torch.yaml: threads_per_node 116 -> 16. The hardcoded 116 oversubscribes cores ~6x and OOMs dataset caching on unified-memory hardware (each worker loads the model + the 89MB style JSON). - scripts/training/*.sh: CUDA_VISIBLE_DEVICES single-GPU for one-GPU machines.
…RScheduler break Two issues that block training on a single Blackwell (sm_121) GPU: - scripts/training/*.sh: add trainer.params.strategy=auto. The default strategy (ddp) initialises NCCL even for one device, and NCCL 2.26.3 hits a shared-mem mismatch on sm_121 (ncclMaxSharedMem exceeds device maxSharedMem). 'auto' uses the single-device strategy and skips NCCL entirely. - WarmupCosLR: stop forwarding the 'verbose' arg to LRScheduler.__init__, which PyTorch >=2.4 removed (TypeError: takes 2 to 3 positional arguments but 4 given).
Prepares the deployment for upstream contribution without regressing multi-GPU
or many-core setups:
- worker.threads_per_node: 16 -> null (auto-detect). Removes the machine-specific
hardcode; restores all-cores behaviour on large servers while staying safe on
fewer-core / unified-memory machines. OOM cap documented as an override.
- scripts/training/*.sh: CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0} so the
single-GPU default is overridable for multi-GPU hosts.
- docs/install_spark.md: full ARM64/Blackwell setup (build, data, cache, train)
with the Spark-specific memory/GPU/ownership notes.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds an optional Docker path to run StyleDrive on NVIDIA DGX Spark (GB10, Grace-Blackwell sm_121, aarch64), which the pinned x86_64/CUDA-11 stack can't support. Additive (new Dockerfile, .dockerignore, download/prepare_dataset.sh, docs/install_spark.md) plus fixes that help any modern setup: hydra-core Python-3.11+ import, WarmupCosLR passing the removed verbose arg (breaks on torch ≥2.4), the HuggingFace /blob/→/resolve/ style-JSON URL, dataset meta_datas nesting, and de-hardcoding threads_per_node (was 116). Hardware knobs (threads_per_node, GPU selection, strategy) are auto- or overridable, so multi-GPU and many-core setups remain unchanged.