diff --git a/POSITRONIC_README.md b/POSITRONIC_README.md index 554cb9b..97ac148 100644 --- a/POSITRONIC_README.md +++ b/POSITRONIC_README.md @@ -8,10 +8,13 @@ In src/openpi/training/config.py modify config with new dataset name. ```python TrainConfig( - name="pi0_positronic_lowmem", - # Here is an example of loading a pi0-FAST model for LoRA finetuning. - # For setting action_dim, action_horizon, and max_token_len, see the comments above. - model=pi0.Pi0Config(paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora"), + name="pi05_positronic_lowmem", + # Pi05 model with LoRA finetuning for low memory usage on Positronic dataset. + model=pi0.Pi0Config( + pi05=True, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ), data=LeRobotPositronicDataConfig( repo_id="", @@ -25,7 +28,7 @@ In src/openpi/training/config.py modify config with new dataset name. Run script to compute normalization stats for your dataset. ```bash -HF_LEROBOT_HOME="path/to/dataset" uv run scripts/compute_norm_stats.py --config-name pi0_positronic_lowmem +HF_LEROBOT_HOME="path/to/dataset" uv run scripts/compute_norm_stats.py --config-name pi05_positronic_lowmem ``` ### Run train @@ -33,13 +36,98 @@ HF_LEROBOT_HOME="path/to/dataset" uv run scripts/compute_norm_stats.py --config- Train policy. ```bash -HF_LEROBOT_HOME="path/to/dataset" XLA_PYTHON_CLIENT_MEM_FRACTION=0.995 uv run scripts/train.py pi0_positronic_lowmem --exp-name= --overwrite +HF_LEROBOT_HOME="path/to/dataset" XLA_PYTHON_CLIENT_MEM_FRACTION=0.995 uv run scripts/train.py pi05_positronic_lowmem --exp-name= --overwrite +``` + +### Serve policy + +After training you could find weights in `checkpoints/pi05_positronic_lowmem//`. Run policy server with: + +```bash +uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi05_positronic_lowmem --policy.dir checkpoints/pi05_positronic_lowmem//29999/ +``` + +By default, the server is served from port 8000, so if you serve it on Nebius or other cloud provider, you need to have this port open: +```bash +ssh -L 8000:localhost:8000 +``` + +Then on the local machine, you will run the inference script (from `positronic` repository root): +```bash +python -m positronic.run_inference sim_pi0 --output_dir=../datasets/inference/ --show_gui --num_iterations=10 --simulation_time=15 +``` + +## Using Docker + +### Build and Push Docker Image + +Build and push the training image to Nebius Container Registry: + +```bash +cd docker +make push +``` + +This will: +1. Build the Docker image with all source code baked in +2. Tag it with version, git SHA, and `latest` +3. Push to the Nebius registry at `cr.eu-north1.nebius.cloud/e00a0ahqzcp9x0xczz/openpi-training` + +### Run Training in Docker + +All commands can be run using the cloud Docker image. The code is already baked into the image, so no source code mounting is needed: + +```bash +docker run --rm -it --gpus all --pull=always \ + -v /datasets:/datasets \ + -v /outputs:/outputs \ + cr.eu-north1.nebius.cloud/e00a0ahqzcp9x0xczz/openpi-training \ + +``` + +#### Docker Examples + +**Compute norm stats:** +```bash +docker run --rm -it --gpus all --pull=always \ + -v /datasets:/datasets \ + -v /outputs/openpi/assets:/openpi/assets \ + -e HF_LEROBOT_HOME=/datasets/ \ + cr.eu-north1.nebius.cloud/e00a0ahqzcp9x0xczz/openpi-training \ + python -m scripts.compute_norm_stats --config-name pi05_positronic_lowmem +``` + +**Run train:** +```bash +docker run --rm -it --gpus all --pull=always \ + -v /datasets:/datasets \ + -v /outputs/openpi/assets:/openpi/assets \ + -v /outputs/checkpoints:/openpi/checkpoints \ + -e HF_LEROBOT_HOME=/datasets/ \ + cr.eu-north1.nebius.cloud/e00a0ahqzcp9x0xczz/openpi-training \ + python -m scripts.train pi05_positronic_lowmem --exp-name= +``` + +**Serve policy:** +```bash +docker run --rm -it --gpus all --pull=always \ + -v /outputs/openpi/checkpoints:/openpi/checkpoints \ + -p 8000:8000 \ + cr.eu-north1.nebius.cloud/e00a0ahqzcp9x0xczz/openpi-training \ + python -m scripts.serve_policy policy:checkpoint \ + --policy.config=pi05_positronic_lowmem \ + --policy.dir=/openpi/checkpoints/pi05_positronic_lowmem//29999/ ``` -### Serve polciy +### Local Development (Optional) -After training you could find weights in `checkpoints/pi0_positronic_lowmem//`. Run policy server with: +For local development with live code editing, use `Dockerfile.training`: ```bash -uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_positronic_lowmem --policy.dir checkpoints/pi0_positronic_lowmem//29999/ +docker build -f docker/Dockerfile.training -t openpi-training-dev . +docker run --rm -it --gpus all \ + -v $PWD:/openpi \ + -v /datasets:/datasets \ + openpi-training-dev \ + bash -lc 'uv pip install --no-deps -e /openpi && ' ``` \ No newline at end of file diff --git a/docker/Dockerfile.training b/docker/Dockerfile.training new file mode 100644 index 0000000..e33b336 --- /dev/null +++ b/docker/Dockerfile.training @@ -0,0 +1,41 @@ +# Local development Dockerfile with bind-mount support. +# The virtualenv and dependencies are pre-installed in the image, then at runtime you mount +# your local source code over /openpi and reinstall in editable mode for live code changes. +# +# Build: +# docker build -f Dockerfile.training -t openpi-training . +# +# Run: +# docker run --rm -it --gpus all \ +# -v $PWD:/openpi \ +# -v /datasets:/datasets \ +# openpi-training \ +# bash -lc 'python -m pip install --no-deps -e /openpi && python scripts/your_script.py' +# +# Note: After mounting your code, reinstall with 'python -m pip install --no-deps -e /openpi' +# to update package metadata while keeping dependencies from the image. +FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04 +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ +RUN apt-get update && apt-get install -y \ + build-essential \ + curl \ + git \ + ffmpeg \ + libturbojpeg \ + && rm -rf /var/lib/apt/lists/* +ENV UV_LINK_MODE=copy +ENV PIP_DEFAULT_TIMEOUT=600 +ENV UV_HTTP_TIMEOUT=600 +RUN uv venv --python 3.11 +ENV VIRTUAL_ENV=/.venv +ENV PATH=/$VIRTUAL_ENV/bin:$PATH +COPY pyproject.toml uv.lock ./ +# Copy workspace packages needed for dependency resolution +COPY packages ./packages +RUN --mount=type=cache,target=/root/.cache/uv \ + uv sync --frozen --extra dev --no-install-project +COPY . /openpi +WORKDIR /openpi +RUN uv pip install --no-deps . +ENV PYTHONPATH=/openpi +CMD ["bash"] diff --git a/docker/Dockerfile.training.cloud b/docker/Dockerfile.training.cloud new file mode 100644 index 0000000..16a44e1 --- /dev/null +++ b/docker/Dockerfile.training.cloud @@ -0,0 +1,50 @@ +FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04 + +# Copy uv for fast Python package management +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ + +# Install minimal system dependencies +RUN apt-get update && apt-get install -y \ + build-essential \ + curl \ + git \ + ffmpeg \ + libturbojpeg \ + && rm -rf /var/lib/apt/lists/* + +# Configure uv and pip timeouts +ENV UV_LINK_MODE=copy +ENV PIP_DEFAULT_TIMEOUT=600 +ENV UV_HTTP_TIMEOUT=600 + +# Create Python virtual environment +RUN uv venv --python 3.11 + +# Activate venv for docker run +ENV VIRTUAL_ENV=/.venv +ENV PATH=/$VIRTUAL_ENV/bin:$PATH + +# Copy project metadata and lock for reproducible deps +COPY pyproject.toml uv.lock ./ +# Copy workspace packages needed for dependency resolution +COPY packages ./packages +RUN --mount=type=cache,target=/root/.cache/uv \ + uv sync --frozen --extra dev --no-install-project + +# Copy source code into the image +COPY . /openpi + +# Set working directory +WORKDIR /openpi + +# Install the project so runtime metadata (e.g. importlib.metadata) is present. +# For local dev containers that bind-mount the repo, run +# `uv pip install --no-deps -e /openpi` once after start so edits reflect +# immediately while preserving metadata. +RUN uv pip install --no-deps . + +# Set Python path +ENV PYTHONPATH=/openpi + +# Default command +CMD ["bash"] diff --git a/docker/Makefile b/docker/Makefile new file mode 100644 index 0000000..5069503 --- /dev/null +++ b/docker/Makefile @@ -0,0 +1,89 @@ +.PHONY: all build tag push clean prune help + +# Image configuration +IMAGE_NAME := openpi-training + +# Extract version from pyproject.toml +VERSION := $(shell grep '^version = ' ../pyproject.toml | sed 's/version = "\(.*\)"/\1/') + +# Get git commit SHA (short) +GIT_SHA := $(shell git rev-parse --short HEAD 2>/dev/null || echo "unknown") + +# Registry URL - can be overridden via environment variable or command line +# Example: make push REGISTRY_URL=cr.eu-north1.nebius.cloud/e00a0ahqzcp9x0xczz +REGISTRY_URL ?= cr.eu-north1.nebius.cloud/e00a0ahqzcp9x0xczz + +# Image tags +TAG_LATEST := $(REGISTRY_URL)/$(IMAGE_NAME):latest +TAG_VERSION := $(REGISTRY_URL)/$(IMAGE_NAME):v$(VERSION) +TAG_SHA := $(REGISTRY_URL)/$(IMAGE_NAME):$(GIT_SHA) +LOCAL_TAG := $(IMAGE_NAME):local + +help: + @echo "OpenPI Training Container - Makefile" + @echo "" + @echo "Configuration:" + @echo " Image Name: $(IMAGE_NAME)" + @echo " Version: $(VERSION)" + @echo " Git SHA: $(GIT_SHA)" + @echo " Registry URL: $(REGISTRY_URL)" + @echo "" + @echo "Targets:" + @echo " make build Build the training image with baked source code" + @echo " make tag Tag the image (depends on build)" + @echo " make push Push all tags to registry (depends on tag)" + @echo " make all Build, tag, and push (alias for push)" + @echo " make clean Remove local training images" + @echo " make prune Remove dangling/unused Docker images" + @echo " make help Show this help message" + @echo "" + @echo "To override registry URL:" + @echo " make push REGISTRY_URL=your-registry-url" + +build: + @echo "Building $(IMAGE_NAME) with source code baked in..." + @if [ -z "$(VERSION)" ]; then \ + echo "Error: Could not extract version from pyproject.toml"; \ + exit 1; \ + fi + docker build -f Dockerfile.training.cloud -t $(LOCAL_TAG) .. + +tag: build + @echo "Tagging image with multiple tags..." + @if [ -z "$(REGISTRY_URL)" ]; then \ + echo "Error: REGISTRY_URL is not set."; \ + echo "Please set it via: make push REGISTRY_URL=your-registry-url"; \ + exit 1; \ + fi + docker tag $(LOCAL_TAG) $(TAG_LATEST) + docker tag $(LOCAL_TAG) $(TAG_VERSION) + docker tag $(LOCAL_TAG) $(TAG_SHA) + @echo "Tagged with:" + @echo " - $(TAG_LATEST)" + @echo " - $(TAG_VERSION)" + @echo " - $(TAG_SHA)" + +push: tag + @echo "Pushing images to Container Registry..." + docker push $(TAG_LATEST) + docker push $(TAG_VERSION) + docker push $(TAG_SHA) + @echo "" + @echo "Successfully pushed to Container Registry!" + @echo "Pull with:" + @echo " docker pull $(TAG_LATEST)" + +all: push + +clean: + @echo "Removing local training images..." + -docker rmi $(LOCAL_TAG) + -docker rmi $(TAG_LATEST) + -docker rmi $(TAG_VERSION) + -docker rmi $(TAG_SHA) + @echo "Local images removed." + +prune: + @echo "Pruning dangling and unused Docker images..." + docker image prune -f + @echo "Prune complete." diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py index 5d97906..185456c 100644 --- a/src/openpi/training/config.py +++ b/src/openpi/training/config.py @@ -426,7 +426,7 @@ def create(self, assets_dirs: pathlib.Path, model_config: _model.BaseModelConfig # We return all data transforms for training and inference. No need to change anything here. return dataclasses.replace( - self.create_base_config(assets_dirs), + self.create_base_config(assets_dirs, model_config), repack_transforms=repack_transform, data_transforms=data_transforms, model_transforms=model_transforms, @@ -815,16 +815,16 @@ def __post_init__(self) -> None: name="pi0_positronic_lowmem", # Here is an example of loading a pi0-FAST model for LoRA finetuning. # For setting action_dim, action_horizon, and max_token_len, see the comments above. - model=pi0.Pi0Config(paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora"), + model=pi0_config.Pi0Config(paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora"), data=LeRobotPositronicDataConfig( - repo_id="stack-cubes-pimm-lerobot-new-convert", + repo_id="pimm", base_config=DataConfig(prompt_from_task=True), ), weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"), num_train_steps=30_000, # Again, make sure to match the model config above when extracting the freeze filter # that specifies which parameters should be frozen during LoRA finetuning. - freeze_filter=pi0.Pi0Config( + freeze_filter=pi0_config.Pi0Config( paligemma_variant="gemma_2b_lora", action_expert_variant="gemma_300m_lora" ).get_freeze_filter(), # Turn off EMA for LoRA finetuning. @@ -834,9 +834,9 @@ def __post_init__(self) -> None: name="pi0_positronic", # Here is an example of loading a pi0-FAST model for LoRA finetuning. # For setting action_dim, action_horizon, and max_token_len, see the comments above. - model=pi0.Pi0Config(), + model=pi0_config.Pi0Config(), data=LeRobotPositronicDataConfig( - repo_id="stack-cubes-pimm-lerobot-new-convert", + repo_id="pimm", base_config=DataConfig(prompt_from_task=True), ), weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi0_base/params"), @@ -844,6 +844,28 @@ def __post_init__(self) -> None: # Turn off EMA for LoRA finetuning. ema_decay=None, ), + TrainConfig( + name="pi05_positronic_lowmem", + # Pi05 model with LoRA finetuning for low memory usage on Positronic dataset. + model=pi0_config.Pi0Config( + pi05=True, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ), + data=LeRobotPositronicDataConfig( + repo_id="pimm", + base_config=DataConfig(prompt_from_task=True), + ), + freeze_filter=pi0_config.Pi0Config( + pi05=True, + paligemma_variant="gemma_2b_lora", + action_expert_variant="gemma_300m_lora", + ).get_freeze_filter(), + weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"), + num_train_steps=30_000, + # Turn off EMA for LoRA finetuning. + ema_decay=None, + ), TrainConfig( name="pi05_libero", model=pi0_config.Pi0Config(pi05=True, action_horizon=10, discrete_state_input=False),