Update on "Support fp8 in AOTInductor + support optional<> in C ABI"

int3 · int3 · commit 2b1de23b2936 · 2023-11-08T10:46:05.000-08:00
This was originally ipiszy's PR: #112358 It turns out that we need to add support for optional types in order to support fp8 gemm (i.e. scaled_mm). Since our ABI-stable C interface can't support optional<> directly, I am passing in optional types via pointer instead. `AtenTensorHandle`s are already pointers, so nothing needs to change there. Only value types need to change. We decided on this approach instead of adding an extra `bool` param to the callee because this simplifies things. Having the same number of arguments regardless of whether we are emitting Python / C++ / ABI-compatible C++ makes codegen easier. There are a number of existing ABI-compatible functions that have optional-typed value parameters. Previously, they just assumed they would never be passed a `nullopt` / `None` at runtime. Changing them to use pointer types now would break ABI stability, so I have created an exclude list for those functions. Finally, I think the current implementation is kind of messy, and only works for FallbackKernels, even though technically ExternKernels could also have the same issue. It also doesn't support optional types nested in lists. I've left FIXME comments for both issues. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler Differential Revision: [D51084289](https://our.internmc.facebook.com/intern/diff/D51084289) [ghstack-poisoned]
diff --git a/.ci/docker/ci_commit_pins/triton-rocm.txt b/.ci/docker/ci_commit_pins/triton-rocm.txt
@@ -1 +1 @@
-e8a35b3968780e48df1374482d56cc6cdbb9e351
+34f8189eae57a23cc15b4b4f032fe25757e0db8e
diff --git a/.github/workflows/build-triton-wheel.yml b/.github/workflows/build-triton-wheel.yml
@@ -38,7 +38,7 @@ jobs:
         device: ["cuda", "rocm"]
         include:
           - device: "rocm"
-            rocm_version: "5.7"
+            rocm_version: "5.6"
           - device: "cuda"
             rocm_version: ""
     timeout-minutes: 40
diff --git a/torch/_inductor/codegen/aoti_runtime/interface.cpp b/torch/_inductor/codegen/aoti_runtime/interface.cpp
@@ -166,19 +166,14 @@ AOTIRuntimeError AOTInductorModelCreate(
       auto constant_map = std::make_shared<torch::aot_inductor::ConstantMap>();
       auto input_map = reinterpret_cast<std::unordered_map<std::string, AtenTensorHandle>*>(constant_map_handle);
 
+      for (auto const& kv : *input_map) {
+        constant_map->emplace(kv.first, kv.second);
+      }
+
       auto model = new torch::aot_inductor::AOTInductorModel(
           constant_map,
           ""
       );
-
-      if (input_map) {
-        for (auto const& kv : *input_map) {
-          constant_map->emplace(kv.first, kv.second);
-        }
-      } else {
-        model->load_constants(/*is_cpu*/true);
-      }
-
       *model_handle = reinterpret_cast<AOTInductorModelHandle>(model);
   })
 }
diff --git a/torch/csrc/inductor/aoti_runtime/model.h b/torch/csrc/inductor/aoti_runtime/model.h
@@ -31,35 +31,6 @@
 #define AOTI_NOINLINE
 #endif
 
-// At codegen time, we write out a binary file called constants.bin.
-// We then turn the raw binary to an object file that exposes this
-// symbol and link it into the final .so.
-// For information on the binary format, see `man objcopy`, under
-// the "binary-architecture" flag:
-// https://man7.org/linux/man-pages/man1/objcopy.1.html
-// todo: use #embed in C++ 23 once available
-extern const uint8_t _binary_constants_bin_start[];
-extern const uint8_t _binary_constants_bin_end[];
-
-#define AOTI_CONST_GPU_ALIGNMENT 64
-
-namespace {
-
-#ifdef USE_CUDA
-
-using CUDAPtr = std::unique_ptr<void, std::function<void(void*)>>;
-
-CUDAPtr RAII_cudaMalloc(size_t num_bytes) {
-  void* data_ptr;
-  AOTI_RUNTIME_DEVICE_CHECK(cudaMalloc((void**)&data_ptr, num_bytes));
-  auto deleter = [](void* ptr) { AOTI_RUNTIME_DEVICE_CHECK(cudaFree(ptr)); };
-  return CUDAPtr(data_ptr, deleter);
-}
-
-#endif // USE_CUDA
-
-} // anonymous namespace
-
 AOTI_NOINLINE static void throw_exception(
     const char* call,
     const char* file,
@@ -228,100 +199,6 @@ class AOTInductorModelBase {
 #endif // USE_CUDA
   }
 
-  void load_constants(bool is_cpu) {
-    size_t num_constants = this->num_constants();
-    constants_map_->reserve(num_constants);
-
-    std::vector<size_t> constants_internal_offset(num_constants);
-    if (!is_cpu) {
-      make_cuda_constant_blob(constants_internal_offset);
-    }
-
-    size_t bytes_read = 0;
-    for (size_t i = 0; i < num_constants; i++) {
-      std::string name = this->constant_name(i);
-      size_t data_size = this->constant_data_size(i);
-      uint8_t* internal_ptr =
-          constant_ptr(constants_internal_offset[i], bytes_read, data_size);
-      bytes_read += data_size;
-
-      // Create at::Tensor from copied memory.
-      auto dtype = this->constant_type(i);
-      auto ndim = this->constant_ndim(i);
-      auto size = this->constant_shape(i);
-      auto stride = this->constant_stride(i);
-      auto offset = this->constant_offset(i);
-
-      auto device_type = aoti_torch_device_type_cuda();
-      if (is_cpu) {
-        device_type = aoti_torch_device_type_cpu();
-      }
-
-      AtenTensorHandle tensor_handle;
-      int device_idx = -1; // should be the same as was used for constant_blob_
-#ifdef USE_CUDA
-      AOTI_RUNTIME_DEVICE_CHECK(cudaGetDevice(&device_idx));
-#endif // USE_CUDA
-      AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob(
-          internal_ptr,
-          ndim,
-          size,
-          stride,
-          offset,
-          dtype,
-          device_type,
-          device_idx,
-          &tensor_handle));
-      constants_map_->emplace(std::move(name), tensor_handle);
-    }
-    this->update_constants_map(constants_map_);
-  }
-
-#ifdef USE_CUDA
-  CUDAPtr&& release_constant_blob() {
-    return std::move(constant_blob_);
-  }
-#endif
-
-  uint8_t* constant_ptr(
-      size_t constant_offset,
-      size_t bytes_read,
-      size_t data_size) {
-#ifdef USE_CUDA
-    auto* constants_ptr = static_cast<uint8_t*>(constant_blob_.get());
-    uint8_t* internal_ptr = constants_ptr + constant_offset;
-    // Copy data to GPU memory
-    // TODO: Handle shared storage case.
-    AOTI_RUNTIME_DEVICE_CHECK(cudaMemcpy(
-        internal_ptr,
-        _binary_constants_bin_start + bytes_read,
-        data_size,
-        cudaMemcpyHostToDevice));
-    return internal_ptr;
-#else // !USE_CUDA
-    // get pointer to constant which is packed in model during compile time.
-    return const_cast<uint8_t*>(_binary_constants_bin_start) + bytes_read;
-#endif // USE_CUDA
-  }
-
-  void make_cuda_constant_blob(std::vector<size_t>& constants_internal_offset) {
-#ifdef USE_CUDA
-    size_t num_constants = this->num_constants();
-    // Compute required blob size with 64-alignment if on GPU.
-    size_t max_blob = 0;
-    for (size_t i = 0; i < num_constants; i++) {
-      size_t data_size = this->constant_data_size(i);
-      if (data_size % AOTI_CONST_GPU_ALIGNMENT) {
-        data_size = AOTI_CONST_GPU_ALIGNMENT +
-            (data_size / AOTI_CONST_GPU_ALIGNMENT) * AOTI_CONST_GPU_ALIGNMENT;
-      }
-      constants_internal_offset[i] = max_blob;
-      max_blob += data_size;
-    }
-    constant_blob_ = RAII_cudaMalloc(max_blob);
-#endif // USE_CUDA
-  }
-
   size_t num_inputs() const {
     return inputs_info_.size();
   }
@@ -450,11 +327,6 @@ class AOTInductorModelBase {
   std::shared_ptr<ConstantMap> constants_map_;
   std::vector<AtenTensorHandle> constants_;
 
-#ifdef USE_CUDA
-  // Holds the blob storage for constants' at::Tensor for CUDA.
-  CUDAPtr constant_blob_;
-#endif // USE_CUDA
-
   // A directory with CUDA binary files, e.g. compiled kernels, etc.
   const std::optional<std::string> cubin_dir_;
 
diff --git a/torch/csrc/inductor/aoti_runtime/model_container.h b/torch/csrc/inductor/aoti_runtime/model_container.h
@@ -12,6 +12,35 @@
 // applies to other files under torch/csrc/inductor/aoti_runtime/.
 #include <torch/csrc/inductor/aoti_runtime/model.h>
 
+// At codegen time, we write out a binary file called constants.bin.
+// We then turn the raw binary to an object file that exposes this
+// symbol and link it into the final .so.
+// For information on the binary format, see `man objcopy`, under
+// the "binary-architecture" flag:
+// https://man7.org/linux/man-pages/man1/objcopy.1.html
+// todo: use #embed in C++ 23 once available
+extern const uint8_t _binary_constants_bin_start[];
+extern const uint8_t _binary_constants_bin_end[];
+
+#define AOTI_CONST_GPU_ALIGNMENT 64
+
+namespace {
+
+#ifdef USE_CUDA
+
+using CUDAPtr = std::unique_ptr<void, std::function<void(void*)>>;
+
+CUDAPtr RAII_cudaMalloc(size_t num_bytes) {
+  void* data_ptr;
+  AOTI_RUNTIME_DEVICE_CHECK(cudaMalloc((void**)&data_ptr, num_bytes));
+  auto deleter = [](void* ptr) { AOTI_RUNTIME_DEVICE_CHECK(cudaFree(ptr)); };
+  return CUDAPtr(data_ptr, deleter);
+}
+
+#endif // USE_CUDA
+
+} // anonymous namespace
+
 namespace torch {
 namespace aot_inductor {
 
@@ -50,10 +79,51 @@ class AOTInductorModelContainer {
       output_names_.push_back(model->output_name(i));
     }
 
-    model->load_constants(is_cpu);
+    size_t num_constants = model->num_constants();
+    constants_->reserve(num_constants);
+
+    std::vector<size_t> constants_internal_offset(num_constants);
+    if (!is_cpu) {
+      make_cuda_constant_blob(model, constants_internal_offset);
+    }
+
+    size_t bytes_read = 0;
+    for (size_t i = 0; i < num_constants; i++) {
+      std::string name = model->constant_name(i);
+      size_t data_size = model->constant_data_size(i);
+      uint8_t* internal_ptr =
+          constant_ptr(constants_internal_offset[i], bytes_read, data_size);
+      bytes_read += data_size;
+
+      // Create at::Tensor from copied memory.
+      auto dtype = model->constant_type(i);
+      auto ndim = model->constant_ndim(i);
+      auto size = model->constant_shape(i);
+      auto stride = model->constant_stride(i);
+      auto offset = model->constant_offset(i);
+
+      auto device_type = aoti_torch_device_type_cuda();
+      if (is_cpu) {
+        device_type = aoti_torch_device_type_cpu();
+      }
+
+      AtenTensorHandle tensor_handle;
+      int device_idx = -1; // should be the same as was used for constant_blob_
 #ifdef USE_CUDA
-    constant_blob_ = model->release_constant_blob();
-#endif
+      AOTI_RUNTIME_DEVICE_CHECK(cudaGetDevice(&device_idx));
+#endif // USE_CUDA
+      AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_create_tensor_from_blob(
+          internal_ptr,
+          ndim,
+          size,
+          stride,
+          offset,
+          dtype,
+          device_type,
+          device_idx,
+          &tensor_handle));
+      constants_->emplace(std::move(name), tensor_handle);
+    }
 
     for (auto& model : models_) {
       model->update_constants_map(constants_);
@@ -194,6 +264,47 @@ class AOTInductorModelContainer {
     lk.lock();
     available_models_.push_back(model);
   }
+
+  uint8_t* constant_ptr(
+      size_t constant_offset,
+      size_t bytes_read,
+      size_t data_size) {
+#ifdef USE_CUDA
+    auto* constants_ptr = static_cast<uint8_t*>(constant_blob_.get());
+    uint8_t* internal_ptr = constants_ptr + constant_offset;
+    // Copy data to GPU memory
+    // TODO: Handle shared storage case.
+    AOTI_RUNTIME_DEVICE_CHECK(cudaMemcpy(
+        internal_ptr,
+        _binary_constants_bin_start + bytes_read,
+        data_size,
+        cudaMemcpyHostToDevice));
+    return internal_ptr;
+#else // !USE_CUDA
+    // get pointer to constant which is packed in model during compile time.
+    return const_cast<uint8_t*>(_binary_constants_bin_start) + bytes_read;
+#endif // USE_CUDA
+  }
+
+  void make_cuda_constant_blob(
+      AOTInductorModel* model,
+      std::vector<size_t>& constants_internal_offset) {
+#ifdef USE_CUDA
+    size_t num_constants = model->num_constants();
+    // Compute required blob size with 64-alignment if on GPU.
+    size_t max_blob = 0;
+    for (size_t i = 0; i < num_constants; i++) {
+      size_t data_size = model->constant_data_size(i);
+      if (data_size % AOTI_CONST_GPU_ALIGNMENT) {
+        data_size = AOTI_CONST_GPU_ALIGNMENT +
+            (data_size / AOTI_CONST_GPU_ALIGNMENT) * AOTI_CONST_GPU_ALIGNMENT;
+      }
+      constants_internal_offset[i] = max_blob;
+      max_blob += data_size;
+    }
+    constant_blob_ = RAII_cudaMalloc(max_blob);
+#endif // USE_CUDA
+  }
 };
 
 } // namespace aot_inductor
diff --git a/torch/testing/_internal/common_fsdp.py b/torch/testing/_internal/common_fsdp.py
@@ -873,7 +873,7 @@ def setUp(self):
 
     @property
     def world_size(self):
-        return min(torch.cuda.device_count(), 8) if torch.cuda.is_available() else 4
+        return torch.cuda.device_count() if torch.cuda.is_available() else 4
 
     @property
     def process_group(self):
