bd8436bc6e
* feat: add cut_cross_entropy * fix: add to input * fix: remove from setup.py * feat: refactor into an integration * chore: ignore lint * feat: add test for cce * fix: set max_steps for liger test * chore: Update base model following suggestion Co-authored-by: Wing Lian <wing.lian@gmail.com> * chore: update special_tokens following suggestion Co-authored-by: Wing Lian <wing.lian@gmail.com> * chore: remove with_temp_dir following comments * fix: plugins aren't loaded * chore: update quotes in error message * chore: lint * chore: lint * feat: enable FA on test * chore: refactor get_pytorch_version * fix: lock cce commit version * fix: remove subclassing UT * fix: downcast even if not using FA and config check * feat: add test to check different attentions * feat: add install to CI * chore: refactor to use parametrize for attention * fix: pytest not detecting test * feat: handle torch lower than 2.4 * fix args/kwargs to match docs * use release version cut-cross-entropy==24.11.4 * fix quotes * fix: use named params for clarity for modal builder * fix: handle install from pip * fix: test check only top level module install * fix: re-add import check * uninstall existing version if no transformers submodule in cce * more dataset fixtures into the cache --------- Co-authored-by: Wing Lian <wing.lian@gmail.com> Co-authored-by: Wing Lian <wing@axolotl.ai>
80 lines
2.0 KiB
Python
80 lines
2.0 KiB
Python
"""
|
|
modal application to run axolotl gpu tests in Modal
|
|
"""
|
|
# pylint: disable=duplicate-code
|
|
|
|
import os
|
|
import pathlib
|
|
import tempfile
|
|
|
|
import jinja2
|
|
import modal
|
|
from jinja2 import select_autoescape
|
|
from modal import App, Image
|
|
|
|
cicd_path = pathlib.Path(__file__).parent.resolve()
|
|
|
|
template_loader = jinja2.FileSystemLoader(searchpath=cicd_path)
|
|
template_env = jinja2.Environment(
|
|
loader=template_loader, autoescape=select_autoescape()
|
|
)
|
|
df_template = template_env.get_template("Dockerfile.jinja")
|
|
|
|
df_args = {
|
|
"AXOLOTL_EXTRAS": os.environ.get("AXOLOTL_EXTRAS", ""),
|
|
"AXOLOTL_ARGS": os.environ.get("AXOLOTL_ARGS", ""),
|
|
"PYTORCH_VERSION": os.environ.get("PYTORCH_VERSION", "2.3.1"),
|
|
"BASE_TAG": os.environ.get("BASE_TAG", "main-base-py3.11-cu121-2.3.1"),
|
|
"CUDA": os.environ.get("CUDA", "121"),
|
|
"GITHUB_REF": os.environ.get("GITHUB_REF", "refs/heads/main"),
|
|
"GITHUB_SHA": os.environ.get("GITHUB_SHA", ""),
|
|
"NIGHTLY_BUILD": os.environ.get("NIGHTLY_BUILD", ""),
|
|
}
|
|
|
|
dockerfile_contents = df_template.render(**df_args)
|
|
|
|
temp_dir = tempfile.mkdtemp()
|
|
with open(pathlib.Path(temp_dir) / "Dockerfile", "w", encoding="utf-8") as f:
|
|
f.write(dockerfile_contents)
|
|
|
|
cicd_image = (
|
|
Image.from_dockerfile(
|
|
pathlib.Path(temp_dir) / "Dockerfile",
|
|
context_mount=None,
|
|
force_build=True,
|
|
gpu="A10G",
|
|
)
|
|
.env(df_args)
|
|
.pip_install("fastapi==0.110.0", "pydantic==2.6.3")
|
|
)
|
|
|
|
app = App("Axolotl CI/CD", secrets=[])
|
|
|
|
|
|
N_GPUS = int(os.environ.get("N_GPUS", 1))
|
|
GPU_CONFIG = modal.gpu.A10G(count=N_GPUS)
|
|
|
|
|
|
def run_cmd(cmd: str, run_folder: str):
|
|
import subprocess # nosec
|
|
|
|
# Propagate errors from subprocess.
|
|
if exit_code := subprocess.call(cmd.split(), cwd=run_folder): # nosec
|
|
exit(exit_code) # pylint: disable=consider-using-sys-exit
|
|
|
|
|
|
@app.function(
|
|
image=cicd_image,
|
|
gpu=GPU_CONFIG,
|
|
timeout=60 * 60,
|
|
cpu=8.0,
|
|
memory=131072,
|
|
)
|
|
def cicd_pytest():
|
|
run_cmd("./cicd/cicd.sh", "/workspace/axolotl")
|
|
|
|
|
|
@app.local_entrypoint()
|
|
def main():
|
|
cicd_pytest.remote()
|