Adds initial CI suport

.github/workflows/unit_test.yaml

@@ -0,0 +1,43 @@
+name: Unit Test
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+
+concurrency:
+  group: unit-test${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
+  cancel-in-progress: true
+
+defaults:
+  run:
+    shell: bash -l -eo pipefail {0}
+
+jobs:
+  unit_tests:
+    if: github.repository_owner == 'pytorch'
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.9', '3.10', '3.11']
+    steps:
+      - name: Check out repo
+        uses: actions/checkout@v4
+      - name: Setup conda env
+        uses: conda-incubator/setup-miniconda@v2
+        with:
+          auto-update-conda: true
+          miniconda-version: "latest"
+          activate-environment: test
+          python-version: ${{ matrix.python-version }}
+      - name: Update pip
+        run: python -m pip install --upgrade pip
+      - name: Install dependencies
+        run: |
+          python -m pip install torch torchvision torchao
+          python -m pip install -e ".[dev]"
+          python -m pip install -e ".[oss]"
+      - name: Run unit tests with coverage
+        run: pytest tests --cov=. --cov-report=xml --durations=20 -vv
+      - name: Upload Coverage to Codecov
+        uses: codecov/codecov-action@v3

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     # PyTorch
     "torchdata>=0.8.0",
     "torchtitan",
+    "torchao",
     # vLLM
     # TODO: pin specific vllm version
     #"vllm==0.10.0",
@@ -36,6 +37,7 @@ dev = [
     "pytest-cov",
     "tensorboard",
     "tomli>=1.1.0",
+    "tomli_w",
     "anyio",
     "pytest-asyncio",
 ]

scripts/common-setup.sh

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.

src/forge/actors/__init__.py

@@ -5,6 +5,5 @@
 # LICENSE file in the root directory of this source tree.
 
 from .collector import Collector
-from .policy import Policy, PolicyRouter
 
-__all__ = ["Collector", "Policy", "PolicyRouter"]
+__all__ = ["Collector"]

src/forge/data/datasets/__init__.py

@@ -4,16 +4,17 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .dataset import DatasetInfo, InfiniteTuneIterableDataset
+from .dataset import DatasetInfo, InfiniteTuneIterableDataset, InterleavedDataset
 from .hf_dataset import HfIterableDataset
 from .packed import PackedDataset
-from .sft_dataset import SFTOutputTransform, sft_iterable_dataset
+from .sft_dataset import sft_iterable_dataset, SFTOutputTransform
 
 __all__ = [
     "DatasetInfo",
     "HfIterableDataset",
+    "InterleavedDataset",
     "InfiniteTuneIterableDataset",
     "PackedDataset",
     "SFTOutputTransform",
     "sft_iterable_dataset",
-]
+]