Fix/causal conv1d cuda128

.gemini/config.yaml

@@ -0,0 +1,12 @@
+have_fun: false
+code_review:
+  disable: false
+  comment_severity_threshold: MEDIUM
+  max_review_comments: -1
+  pull_request_opened:
+    help: false
+    # disable PR summaries
+    summary: false
+    code_review: true
+    include_drafts: false
+ignore_patterns: []

.github/workflows/cpu_ci.yaml

@@ -0,0 +1,89 @@
+name: SkyRL
+
+on: 
+  push: 
+    branches: 
+      - main 
+      - rc/**
+  pull_request:
+
+permissions:
+  checks: write   # for status checks to appear
+  contents: read
+
+# Cancel runs for previous commits on the same branch
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  check_code_quality:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        # This is the version of the action for setting up Python, not the Python version.
+        uses: actions/setup-python@v5
+        with:
+          # Semantic version range syntax or exact version of a Python version
+          python-version: '3.12'
+          cache: 'pip'
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          activate-environment: true
+      - name: Run pre-commit hooks
+        run: uv pip install pre-commit; pre-commit run --all-files --config .pre-commit-config.yaml
+
+  skyrl_tests:
+    needs: check_code_quality
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash
+        working-directory: ./skyrl-train
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        # This is the version of the action for setting up Python, not the Python version.
+        uses: actions/setup-python@v5
+        with:
+          # Semantic version range syntax or exact version of a Python version
+          python-version: '3.12'
+          cache: 'pip'
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          activate-environment: true
+      - name: Install skyrl-train
+        run:  uv sync --frozen --extra dev # installs from lock file
+      - name: Run cpu tests
+        run:  uv run --frozen pytest tests/cpu/
+
+  skyrl_gym_tests:
+    needs: check_code_quality
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash
+        working-directory: ./skyrl-gym
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        # This is the version of the action for setting up Python, not the Python version.
+        uses: actions/setup-python@v5
+        with:
+          # Semantic version range syntax or exact version of a Python version
+          python-version: '3.12'
+          cache: 'pip'
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          activate-environment: true
+      - name: Install skyrl-gym
+        run:  uv sync --frozen --extra dev # installs from lock file
+      - name: Run cpu tests
+        run:  uv run --frozen pytest tests/

.github/workflows/cpu_skyrl_tx.yaml

@@ -0,0 +1,67 @@
+name: SkyRL-tx-CPU
+
+on:
+  push:
+    branches: [ main ]
+    paths:
+      - 'skyrl-tx/**'
+      - '.github/workflows/cpu_skyrl_tx.yaml'
+  pull_request:
+    paths:
+      - 'skyrl-tx/**'
+      - '.github/workflows/cpu_skyrl_tx.yaml'
+  workflow_dispatch:
+
+permissions:
+  checks: write   # for status checks to appear
+  contents: read
+
+# Cancel runs for previous commits on the same branch
+concurrency:
+  group: skyrl-tx-${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  skyrl_tx_tests:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash
+        working-directory: ./skyrl-tx
+    steps:
+    - uses: actions/checkout@v4
+    - name: Install uv
+      run: |
+        curl -LsSf https://astral.sh/uv/install.sh | sh
+    # - name: Check if reference docs are up to date
+    #   run: |
+    #     uv run --extra dev typer tx/run/main.py utils docs --name tx --output docs/reference.md && git diff --exit-code docs/reference.md
+    # - name: Test docs
+    #   run: |
+    #     uv run --extra dev mkdocs build --strict
+    - name: Run lint
+      run: |
+        uvx ruff check
+    # - name: Run type checking
+    #   run: |
+    #     uv run --extra tinker --extra dev ty check
+    - name: Run pytest
+      run: |
+        uv run --extra tinker --extra dev pytest --forked -s tests
+    - name: Run a single training step
+      run: |
+        uv run tx train --model pcmoritz/qwen3-tiny-test --dataset mahiatlinux/TinyStories-GPT4-V2-50K-SUBSET --output-dir /tmp --batch-size 2 --max-steps 1 --optimizer-args '{"learning_rate": 0.002, "weight_decay": 0.1}'
+        uv run tx train --model pcmoritz/qwen3-tiny-test --dataset mahiatlinux/TinyStories-GPT4-V2-50K-SUBSET --output-dir /tmp --batch-size 2 --max-steps 1 --optimizer-args '{"learning_rate": 0.002, "weight_decay": 0.1}' --tp-size 2
+    - name: Run a single fine-tuning step on a chat dataset
+      run: |
+        uv run --with jinja2 tx train --model pcmoritz/qwen3-tiny-test --dataset smangrul/ultrachat-feedback-10k-chatml --output-dir /tmp --batch-size 2 --max-steps 1 --loader tx.loaders.chat --load-checkpoint-path /tmp
+    - name: Run a single fine-tuning step with Qwen3 MoE
+      run: |
+        uv run --with huggingface_hub hf download trl-internal-testing/tiny-Qwen3MoeForCausalLM --local-dir /tmp/qwen3_moe
+        uv run --with jinja2 tx train --model trl-internal-testing/tiny-Qwen3MoeForCausalLM --dataset smangrul/ultrachat-feedback-10k-chatml --output-dir /tmp --batch-size 2 --max-steps 1 --loader tx.loaders.chat --load-checkpoint-path /tmp/qwen3_moe
+    - name: Test experiment tracker integration
+      run: |
+        WANDB_MODE=offline WANDB_API_KEY=dummy uv run --with wandb tx train --model pcmoritz/qwen3-tiny-test --dataset mahiatlinux/TinyStories-GPT4-V2-50K-SUBSET --output-dir /tmp --batch-size 2 --max-steps 1 --tracker wandb --tracker-args '{"name": "Qwen3-8B", "project": "tx"}'
+    - name: Run engine benchmarks
+      run: |
+        uv run --extra tinker --extra dev python benchmarks/benchmark_engine.py --base-model trl-internal-testing/tiny-Qwen3ForCausalLM --backend-config '{"max_lora_adapters": 3, "max_lora_rank": 1}' --num-warmup-steps 1 --num-steps 1 --num-requests 1 --seq-len 8 --sample-max-tokens 16

.github/workflows/gpu_ci.yaml

@@ -0,0 +1,50 @@
+name: SkyRL-GPU
+
+on: 
+  push: 
+    branches: 
+      - main
+    paths:
+      - 'skyrl-train/**'
+      - '!skyrl-train/docs/**'
+      - '!skyrl-train/examples/**'
+      - '.github/workflows/**'
+  workflow_dispatch:
+
+
+permissions:
+  checks: write   # for status checks to appear
+  contents: read
+
+jobs:
+
+  skyrl_tests:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash
+        working-directory: ./skyrl-train
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        # This is the version of the action for setting up Python, not the Python version.
+        uses: actions/setup-python@v5
+        with:
+          # Semantic version range syntax or exact version of a Python version
+          python-version: '3.12'
+          cache: 'pip'
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          activate-environment: true
+      - name: Install basic dependencies
+        run: uv pip install anyscale==0.24.79 typer==0.9.0
+      # Run tests
+      - name: GPU tests
+        env:
+          ANYSCALE_CLI_TOKEN: ${{ secrets.ANYSCALE_CLI_TOKEN }}
+          ANYSCALE_HOST: https://console.anyscale.com
+        run: |
+          anyscale job submit -f ci/anyscale_gpu_ci.yaml --timeout 10000
+          anyscale job wait --cloud sky-anyscale-aws-us-east-1 --name skyrl-train-gpu-ci --timeout 10000

.github/workflows/gpu_ci_megatron.yaml

@@ -0,0 +1,55 @@
+name: SkyRL-GPU-Megatron
+
+on: 
+  push: 
+    branches: 
+      - main
+    paths:
+      - 'skyrl-train/skyrl_train/workers/**'
+      - 'skyrl-train/skyrl_train/distributed/megatron/**'
+      - 'skyrl-train/skyrl_train/config/**'
+      - 'skyrl-train/skyrl_train/trainer.py'
+      - 'skyrl-train/tests/gpu/gpu_ci/**'
+      - 'skyrl-train/pyproject.toml'
+      - '!skyrl-train/docs/**'
+      - '!skyrl-train/examples/**'
+      - '.github/workflows/**'
+  workflow_dispatch:
+
+
+permissions:
+  checks: write   # for status checks to appear
+  contents: read
+
+jobs:
+
+  skyrl_tests_megatron:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash
+        working-directory: ./skyrl-train
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        # This is the version of the action for setting up Python, not the Python version.
+        uses: actions/setup-python@v5
+        with:
+          # Semantic version range syntax or exact version of a Python version
+          python-version: '3.12'
+          cache: 'pip'
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          activate-environment: true
+      - name: Install basic dependencies
+        run: uv pip install anyscale==0.24.79 typer==0.9.0
+      # Run tests
+      - name: GPU tests
+        env:
+          ANYSCALE_CLI_TOKEN: ${{ secrets.ANYSCALE_CLI_TOKEN }}
+          ANYSCALE_HOST: https://console.anyscale.com
+        run: |
+          anyscale job submit -f ci/anyscale_gpu_ci_megatron.yaml --timeout 5000
+          anyscale job wait --cloud sky-anyscale-aws-us-east-1 --name skyrl-train-gpu-ci-megatron --timeout 5000

.github/workflows/gpu_e2e_ci.yaml

@@ -0,0 +1,47 @@
+name: SkyRL-GPU-E2E-CI
+
+on:
+  schedule:
+    - cron: '5 8 * * *'   # Every day at 08:05 UTC (~00:05 PST / ~01:05 PDT)
+  workflow_dispatch:
+
+permissions:
+  checks: write   # for status checks to appear
+  contents: read
+
+jobs:
+
+  skyrl_gpu_e2e_test:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash
+        working-directory: ./skyrl-train
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        # This is the version of the action for setting up Python, not the Python version.
+        uses: actions/setup-python@v5
+        with:
+          # Semantic version range syntax or exact version of a Python version
+          python-version: '3.12'
+          cache: 'pip'
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          activate-environment: true
+      - name: Install basic dependencies
+        run: uv pip install anyscale==0.24.79 typer==0.9.0
+      - name: Install envsubst
+        run: sudo apt-get update && sudo apt-get install -y gettext-base
+      - name: Basic convergence test
+        env:
+          ANYSCALE_CLI_TOKEN: ${{ secrets.ANYSCALE_CLI_TOKEN }}
+          ANYSCALE_HOST: https://console.anyscale.com
+          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+        run: |
+          envsubst < ci/anyscale_gpu_e2e_test.yaml > ci/anyscale_gpu_e2e_test_envsubst.yaml
+          anyscale job submit -f ci/anyscale_gpu_e2e_test_envsubst.yaml --timeout 4500
+          anyscale job wait --cloud sky-anyscale-aws-us-east-1 --name skyrl-train-gpu-e2e-test --timeout 4500
+          rm -f ci/anyscale_gpu_e2e_test_envsubst.yaml

.github/workflows/gpu_e2e_ci_megatron.yaml

@@ -0,0 +1,47 @@
+name: SkyRL-GPU-E2E-CI-Megatron
+
+on:
+  schedule:
+    - cron: '5 8 * * *'   # Every day at 08:05 UTC (~00:05 PST / ~01:05 PDT)
+  workflow_dispatch:
+
+permissions:
+  checks: write   # for status checks to appear
+  contents: read
+
+jobs:
+
+  skyrl_gpu_e2e_test_megatron:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash
+        working-directory: ./skyrl-train
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        # This is the version of the action for setting up Python, not the Python version.
+        uses: actions/setup-python@v5
+        with:
+          # Semantic version range syntax or exact version of a Python version
+          python-version: '3.12'
+          cache: 'pip'
+      - name: Install the latest version of uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          activate-environment: true
+      - name: Install basic dependencies
+        run: uv pip install anyscale==0.24.79 typer==0.9.0
+      - name: Install envsubst
+        run: sudo apt-get update && sudo apt-get install -y gettext-base
+      - name: Basic convergence test
+        env:
+          ANYSCALE_CLI_TOKEN: ${{ secrets.ANYSCALE_CLI_TOKEN }}
+          ANYSCALE_HOST: https://console.anyscale.com
+          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
+        run: |
+          envsubst < ci/anyscale_gpu_e2e_test_megatron.yaml > ci/anyscale_gpu_e2e_test_megatron_envsubst.yaml
+          anyscale job submit -f ci/anyscale_gpu_e2e_test_megatron_envsubst.yaml --timeout 4500
+          anyscale job wait --cloud sky-anyscale-aws-us-east-1 --name skyrl-train-gpu-e2e-test-megatron --timeout 4500
+          rm -f ci/anyscale_gpu_e2e_test_megatron_envsubst.yaml