Commit d2dcbe90 authored by Gaetan Lepage's avatar Gaetan Lepage
Browse files
parent 35c1f13b
Loading
Loading
Loading
Loading
+0 −25
Original line number Diff line number Diff line
From 4c2b65c47d328c2f20cc74adcec2286fee6cb5de Mon Sep 17 00:00:00 2001
From: Yaroslav Bolyukin <iam@lach.pw>
Date: Tue, 30 Jan 2024 18:18:35 +0100
Subject: [PATCH] fix: allow building without git

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index e01c008..92eca62 100644
--- a/setup.py
+++ b/setup.py
@@ -71,7 +71,7 @@ def get_flash_version() -> str:
             ["git", "describe", "--tags", "--always"],
             cwd=flash_dir,
         ).decode("ascii")[:-1]
-    except subprocess.CalledProcessError:
+    except Exception:
         version = flash_dir / "version.txt"
         if version.is_file():
             return version.read_text().strip()
-- 
2.43.0
+147 −63
Original line number Diff line number Diff line
@@ -3,75 +3,83 @@
  stdenv,
  buildPythonPackage,
  fetchFromGitHub,
  which,

  # build-system
  torch,
  setuptools,
  # runtime dependencies

  # buildInputs
  openmp,

  # dependencies
  numpy,
  torch,
  # check dependencies
  pytestCheckHook,
  pynvml,

  # tests
  einops,
  fairscale,
  hydra-core,
  networkx,
  pytest-cov-stub,
  # , pytest-mpi
  pytest-timeout,
  # , pytorch-image-models
  hydra-core,
  fairscale,
  pytestCheckHook,
  scipy,
  cmake,
  ninja,
  triton,
  networkx,
  #, apex
  einops,
  transformers,
  timm,
  #, flash-attn
  openmp,
  transformers,
  triton,
  python,

  # passthru
  xformers,
  writableTmpDirAsHomeHook,
}:
let
  inherit (torch) cudaCapabilities cudaPackages cudaSupport;

  # version 0.0.32.post2 was confirmed to break CUDA.
  # Remove this note once the latest published revision "just works".
  version = "0.0.30";
  effectiveStdenv = if cudaSupport then cudaPackages.backendStdenv else stdenv;
in
buildPythonPackage.override { stdenv = effectiveStdenv; } {
buildPythonPackage.override { stdenv = effectiveStdenv; } (finalAttrs: {
  pname = "xformers";
  inherit version;
  version = "0.0.35";
  pyproject = true;

  src = fetchFromGitHub {
    owner = "facebookresearch";
    repo = "xformers";
    tag = "v${version}";
    tag = "v${finalAttrs.version}";
    fetchSubmodules = true;
    hash = "sha256-ozaw9z8qnGpZ28LQNtwmKeVnrn7KDWNeJKtT6g6Q/W0=";
    hash = "sha256-UqpRHLN0INpW6sA8DbQCSeL8uhS+IoW60UPVUIh1NY0=";
  };

  patches = [ ./0001-fix-allow-building-without-git.patch ];
  # ModuleNotFoundError: No module named 'xformers.components'
  postPatch = ''
    touch xformers/components/__init__.py
    touch xformers/components/attention/__init__.py
  '';

  build-system = [ setuptools ];
  build-system = [
    setuptools
    torch
  ];

  preBuild = ''
    cat << EOF > ./xformers/version.py
    # noqa: C801
    __version__ = "${version}"
    EOF

    export MAX_JOBS=$NIX_BUILD_CORES
  '';

  env = lib.attrsets.optionalAttrs cudaSupport {
    TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
    # Don't silently fallback to a non-CUDA build
    FORCE_CUDA = 1;

    TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" cudaCapabilities}";
  };

  buildInputs =
    lib.optional stdenv.hostPlatform.isDarwin openmp
    lib.optionals stdenv.hostPlatform.isDarwin [
      openmp
    ]
    ++ lib.optionals cudaSupport (
      with cudaPackages;
      [
        # flash-attn build
        cuda_cudart # cuda_runtime_api.h
        libcusparse # cusparse.h
        cuda_cccl # nv/target
@@ -81,51 +89,127 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } {
      ]
    );

  nativeBuildInputs = [
    ninja
    which
  nativeBuildInputs =
    lib.optionals cudaSupport [
      cudaPackages.cuda_nvcc
    ]
  ++ lib.optionals cudaSupport (with cudaPackages; [ cuda_nvcc ])
  ++ lib.optional stdenv.hostPlatform.isDarwin openmp.dev;
    ++ lib.optionals stdenv.hostPlatform.isDarwin [
      openmp.dev
    ];

  dependencies = [
    numpy
    torch
  ]
  ++ lib.optionals cudaSupport [
    pynvml
  ];

  pythonImportsCheck = [ "xformers" ];

  # Has broken 0.03 version:
  # https://github.com/NixOS/nixpkgs/pull/285495#issuecomment-1920730720
  passthru.skipBulkUpdate = true;

  dontUseCmakeConfigure = true;

  # see commented out missing packages
  doCheck = false;
  pythonImportsCheck = [
    "xformers"
    "xformers.components"
    "xformers.components.attention"
  ];

  nativeCheckInputs = [
    pytestCheckHook
    einops
    fairscale
    hydra-core
    networkx
    pytest-cov-stub
    pytest-timeout
    hydra-core
    fairscale
    pytestCheckHook
    scipy
    cmake
    networkx
    triton
    # apex
    einops
    transformers
    timm
    # flash-attn
    transformers
    triton
  ];

  preCheck =
    # Otherwise the CPP bindings are not available and the GPU tests fail with:
    # `fa2F@2.5.7-pt` is not supported because:
    #     xFormers wasn't build with CUDA support: False
    ''
      rm -rf xformers
    ''
    # Display information about the installation
    + ''
      ${python.interpreter} -m xformers.info
    '';

  enabledTestPaths = [
    "tests"
  ];

  disabledTestPaths = [
    # Those tests require access to the GPU and should be tagged accordingly:
    # https://github.com/facebookresearch/xformers/pull/1385
    "tests/test_fwbw_overlap.py"
  ];

  disabledTests =
    # The following tests fail without cudaSupport
    lib.optionals (!cudaSupport) [
      # AssertionError: Torch not compiled with CUDA enabled
      "test_flash_gqa_wrong_strides"
      "test_memeff_compile"
      "test_paged_attention"
      "test_paged_attention_flash"
      "test_triton_splitk_decoder"
      "test_triton_splitk_decoder_manyqueries"
      "test_unsupported_dropout_combine_flash_cutlass"

      # AssertionError: Should use Flash-Decoding with BMHK MQA
      "test_dispatch_decoding_bmhk"

      # AssertionError: Should use Flash-Decoding with MQA
      "test_dispatch_decoding_bmghk"
    ];

  passthru.gpuCheck = xformers.overridePythonAttrs (old: {
    requiredSystemFeatures = [ "cuda" ];

    # Run all tests, including the ones that need a GPU
    disabledTestPaths = [ ];

    disabledTests = (old.disabledTests or [ ]) ++ [
      # `fa3F@0.0.0` is not supported because:
      #   operator wasn't built - see `python -m xformers.info` for more info
      "test_merge_training"

      # Not enough GPU memory (on a 20G RTX 4000 SFF Ada)
      # triton.runtime.errors.OutOfResources: out of resource:
      # shared memory, Required: 147456, Hardware limit: 101376.
      "test_consistency"
      "test_forward"
      "test_logsumexp"
      "test_mqa_decoding"
      "test_tree_attention"

      # Tolerance issues:
      # AssertionError: cutlassF-pt+cutlassB-pt:key: out=4.65625 and ref=3.1875
      # (diff=0.251953125 > 0) at (np.int64(0), np.int64(0), np.int64(4))
      # of shape (1, 4, 8) / atol=0.9, rtol=0.1/ total failing elements: 1 (3.12%)
      "test_backward"

      # AssertionError: Legacy CUDA profiling requires use_cpu=True
      "test_profiler_dispatcher_stream_workaround"

      # RuntimeError: two_four_sgemm_cutlass, /build/source/xformers/csrc/sparse24/gemm.cu:190,
      # Got CUTLASS error: Error Internal
      "test_linearw24"
    ];

    nativeCheckInputs = old.nativeCheckInputs ++ [
      writableTmpDirAsHomeHook
    ];
  });

  meta = {
    description = "Collection of composable Transformer building blocks";
    homepage = "https://github.com/facebookresearch/xformers";
    changelog = "https://github.com/facebookresearch/xformers/blob/${version}/CHANGELOG.md";
    changelog = "https://github.com/facebookresearch/xformers/blob/${finalAttrs.src.tag}/CHANGELOG.md";
    license = lib.licenses.bsd3;
    maintainers = with lib.maintainers; [ happysalada ];
  };
}
})