diff --git a/pkgs/development/python-modules/liger-kernel/default.nix b/pkgs/development/python-modules/liger-kernel/default.nix new file mode 100644 index 000000000000..8454e6933a35 --- /dev/null +++ b/pkgs/development/python-modules/liger-kernel/default.nix @@ -0,0 +1,45 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + setuptools, + setuptools-scm, + torch, + triton, +}: + +buildPythonPackage (finalAttrs: { + pname = "liger-kernel"; + version = "0.8.0"; + pyproject = true; + + src = fetchFromGitHub { + owner = "linkedin"; + repo = "liger-kernel"; + tag = "v${finalAttrs.version}"; + hash = "sha256-MAJ0goMxi8qA+ODd/kLIO7hp4dmNIOPLpSdkrm7hnnQ="; + }; + + build-system = [ + setuptools + setuptools-scm + ]; + + dependencies = [ + torch + triton + ]; + + # Requires NVIDIA driver. + doCheck = false; + + pythonImportsCheck = [ "liger_kernel" ]; + + meta = { + description = "Efficient Triton Kernels for LLM Training"; + homepage = "https://github.com/linkedin/liger-kernel"; + changelog = "https://github.com/linkedin/liger-kernel/releases/tag/${finalAttrs.src.tag}"; + license = lib.licenses.bsd3; + maintainers = with lib.maintainers; [ jherland ]; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index fe245b06a37b..ff667c498a02 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -9513,6 +9513,8 @@ self: super: with self; { lifx-emulator-core = callPackage ../development/python-modules/lifx-emulator-core { }; + liger-kernel = callPackage ../development/python-modules/liger-kernel { }; + lightgbm = callPackage ../development/python-modules/lightgbm { }; lightify = callPackage ../development/python-modules/lightify { };