Fix precision overflow problem (#500)

bytedance · Apr 6, 2023 · 7e5bed6 · 7e5bed6
1 parent b538578
commit 7e5bed6
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 8 deletions.
diff --git a/lightseq/csrc/layers_new/includes/sdpa_layer.h b/lightseq/csrc/layers_new/includes/sdpa_layer.h
@@ -26,8 +26,8 @@ class SDPALayer : public Layer {
   int _head_dim;
 
  public:
-  SDPALayer(int max_batch_tokens, int max_seq_len, int head_dim, int num_heads,
-            float attn_prob_dropout_ratio);
+  SDPALayer(size_t max_batch_tokens, size_t max_seq_len, size_t head_dim,
+            size_t num_heads, float attn_prob_dropout_ratio);
 
   virtual ~SDPALayer() {}
 

diff --git a/lightseq/csrc/layers_new/sdpa_layer.cpp b/lightseq/csrc/layers_new/sdpa_layer.cpp
@@ -7,8 +7,8 @@ Scaled Dot Product Attention
 See paper "Attention is all you need" for details"
 */
 template <typename T1, typename T2>
-SDPALayer<T1, T2>::SDPALayer(int max_batch_tokens, int max_seq_len,
-                             int head_dim, int num_heads,
+SDPALayer<T1, T2>::SDPALayer(size_t max_batch_tokens, size_t max_seq_len,
+                             size_t head_dim, size_t num_heads,
                              float attn_prob_dropout_ratio)
     : Layer("SDPALayer"),
       // for training, max_batch_tokens =

diff --git a/lightseq/csrc/lsflow/includes/tensor.h b/lightseq/csrc/lsflow/includes/tensor.h
@@ -34,7 +34,7 @@ class Tensor {
 
   static int global_tensor_id;
   TensorPtr _original_tensor;
-  int _offset = 0;
+  size_t _offset = 0;
 
  public:
   // Applies to tensors using FixedMemory and SharedMemory memory types.
@@ -63,7 +63,7 @@ class Tensor {
   // Set a specific offset value for a tensor whose memory type is OffsetMemory.
   // Note that the `offset` value here represents the number of elements, not
   // bytes.
-  void set_offset(int offset, Shape shape);
+  void set_offset(size_t offset, Shape shape);
 
   // This method executes logic differently in different situations.
   //

diff --git a/lightseq/csrc/lsflow/tensor.cpp b/lightseq/csrc/lsflow/tensor.cpp
@@ -140,7 +140,7 @@ void Tensor::set_tensor(const char* inp) { set_tensor(const_cast<char*>(inp)); }
 
 void Tensor::set_shape(Shape shape) { _shape = shape; }
 
-void Tensor::set_offset(int offset, Shape shape) {
+void Tensor::set_offset(size_t offset, Shape shape) {
   if (_original_tensor == nullptr) {
     printf("Error! tensor %s set_offset without original tensor",
            _name.c_str());
@@ -226,7 +226,7 @@ void Tensor::print_tensor(int size) {
     printf(", tensor dtype: %d", _dtype);
 
     if (_mtype == LSMemoryType::OffsetMemory) {
-      printf(", offset is %d\n", _offset);
+      printf(", offset is %zu\n", _offset);
     } else {
       printf("\n");
     }