mindspore 2.2.10__cp39-cp39-win_amd64.whl → 2.2.14__cp39-cp39-win_amd64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of mindspore might be problematic. Click here for more details.
- mindspore/.commit_id +1 -1
- mindspore/__init__.py +2 -1
- mindspore/_c_dataengine.cp39-win_amd64.pyd +0 -0
- mindspore/_c_expression.cp39-win_amd64.pyd +0 -0
- mindspore/_c_mindrecord.cp39-win_amd64.pyd +0 -0
- mindspore/_extends/parse/__init__.py +3 -2
- mindspore/_extends/parse/parser.py +6 -1
- mindspore/_extends/parse/standard_method.py +12 -2
- mindspore/common/_utils.py +16 -0
- mindspore/common/tensor.py +0 -2
- mindspore/communication/management.py +3 -0
- mindspore/context.py +34 -4
- mindspore/dataset/engine/cache_client.py +8 -5
- mindspore/dataset/engine/datasets.py +23 -0
- mindspore/dataset/engine/validators.py +1 -1
- mindspore/dataset/vision/py_transforms_util.py +2 -2
- mindspore/dnnl.dll +0 -0
- mindspore/experimental/optim/lr_scheduler.py +5 -6
- mindspore/jpeg62.dll +0 -0
- mindspore/mindrecord/tools/cifar100_to_mr.py +49 -57
- mindspore/mindrecord/tools/cifar10_to_mr.py +46 -55
- mindspore/mindrecord/tools/csv_to_mr.py +3 -8
- mindspore/mindrecord/tools/mnist_to_mr.py +4 -9
- mindspore/mindrecord/tools/tfrecord_to_mr.py +1 -4
- mindspore/mindspore_backend.dll +0 -0
- mindspore/mindspore_common.dll +0 -0
- mindspore/mindspore_core.dll +0 -0
- mindspore/mindspore_glog.dll +0 -0
- mindspore/mindspore_shared_lib.dll +0 -0
- mindspore/nn/layer/activation.py +1 -1
- mindspore/nn/layer/embedding.py +2 -2
- mindspore/nn/layer/flash_attention.py +48 -135
- mindspore/nn/loss/loss.py +1 -1
- mindspore/nn/optim/ada_grad.py +2 -2
- mindspore/nn/optim/sgd.py +3 -2
- mindspore/nn/wrap/__init__.py +4 -2
- mindspore/nn/wrap/cell_wrapper.py +6 -3
- mindspore/numpy/math_ops.py +1 -1
- mindspore/opencv_core452.dll +0 -0
- mindspore/opencv_imgcodecs452.dll +0 -0
- mindspore/opencv_imgproc452.dll +0 -0
- mindspore/ops/__init__.py +3 -0
- mindspore/ops/_grad_experimental/grad_array_ops.py +0 -31
- mindspore/ops/_grad_experimental/grad_comm_ops.py +4 -2
- mindspore/ops/_grad_experimental/grad_inner_ops.py +8 -0
- mindspore/ops/_grad_experimental/grad_math_ops.py +37 -17
- mindspore/ops/_op_impl/aicpu/__init__.py +1 -0
- mindspore/ops/_op_impl/aicpu/generate_eod_mask.py +38 -0
- mindspore/ops/_op_impl/aicpu/linear_sum_assignment.py +21 -2
- mindspore/ops/function/array_func.py +6 -5
- mindspore/ops/function/debug_func.py +1 -1
- mindspore/ops/function/linalg_func.py +21 -11
- mindspore/ops/function/math_func.py +3 -0
- mindspore/ops/function/nn_func.py +13 -11
- mindspore/ops/function/parameter_func.py +2 -0
- mindspore/ops/function/sparse_unary_func.py +2 -2
- mindspore/ops/function/vmap_func.py +1 -0
- mindspore/ops/operations/__init__.py +5 -2
- mindspore/ops/operations/_embedding_cache_ops.py +1 -1
- mindspore/ops/operations/_grad_ops.py +3 -4
- mindspore/ops/operations/_inner_ops.py +56 -1
- mindspore/ops/operations/_quant_ops.py +4 -4
- mindspore/ops/operations/_rl_inner_ops.py +1 -1
- mindspore/ops/operations/array_ops.py +15 -4
- mindspore/ops/operations/custom_ops.py +1 -1
- mindspore/ops/operations/debug_ops.py +1 -1
- mindspore/ops/operations/image_ops.py +3 -3
- mindspore/ops/operations/inner_ops.py +49 -0
- mindspore/ops/operations/math_ops.py +65 -3
- mindspore/ops/operations/nn_ops.py +95 -28
- mindspore/ops/operations/random_ops.py +2 -0
- mindspore/ops/operations/sparse_ops.py +4 -4
- mindspore/ops/silent_check.py +162 -0
- mindspore/parallel/__init__.py +3 -2
- mindspore/parallel/_auto_parallel_context.py +82 -3
- mindspore/parallel/_parallel_serialization.py +34 -2
- mindspore/parallel/_tensor.py +3 -1
- mindspore/parallel/_transformer/transformer.py +8 -8
- mindspore/parallel/checkpoint_transform.py +191 -45
- mindspore/profiler/parser/ascend_cluster_generator.py +111 -0
- mindspore/profiler/parser/ascend_communicate_generator.py +315 -0
- mindspore/profiler/parser/ascend_flops_generator.py +8 -2
- mindspore/profiler/parser/ascend_fpbp_generator.py +8 -2
- mindspore/profiler/parser/ascend_hccl_generator.py +2 -2
- mindspore/profiler/parser/ascend_msprof_exporter.py +30 -6
- mindspore/profiler/parser/ascend_msprof_generator.py +16 -5
- mindspore/profiler/parser/ascend_op_generator.py +15 -7
- mindspore/profiler/parser/ascend_timeline_generator.py +5 -2
- mindspore/profiler/parser/base_timeline_generator.py +11 -3
- mindspore/profiler/parser/cpu_gpu_timeline_generator.py +2 -1
- mindspore/profiler/parser/framework_parser.py +8 -2
- mindspore/profiler/parser/memory_usage_parser.py +8 -2
- mindspore/profiler/parser/minddata_analyzer.py +8 -2
- mindspore/profiler/parser/minddata_parser.py +1 -1
- mindspore/profiler/parser/msadvisor_analyzer.py +4 -2
- mindspore/profiler/parser/msadvisor_parser.py +9 -3
- mindspore/profiler/profiling.py +97 -25
- mindspore/rewrite/api/node.py +1 -1
- mindspore/rewrite/api/symbol_tree.py +2 -2
- mindspore/rewrite/parsers/for_parser.py +6 -6
- mindspore/rewrite/parsers/module_parser.py +4 -4
- mindspore/tinyxml2.dll +0 -0
- mindspore/train/callback/_checkpoint.py +8 -8
- mindspore/train/callback/_landscape.py +2 -3
- mindspore/train/callback/_summary_collector.py +6 -7
- mindspore/train/dataset_helper.py +6 -0
- mindspore/train/model.py +17 -5
- mindspore/train/serialization.py +6 -1
- mindspore/train/summary/_writer_pool.py +1 -1
- mindspore/train/summary/summary_record.py +5 -6
- mindspore/turbojpeg.dll +0 -0
- mindspore/version.py +1 -1
- {mindspore-2.2.10.dist-info → mindspore-2.2.14.dist-info}/METADATA +3 -2
- {mindspore-2.2.10.dist-info → mindspore-2.2.14.dist-info}/RECORD +117 -124
- mindspore/ops/_op_impl/_custom_op/flash_attention/__init__.py +0 -0
- mindspore/ops/_op_impl/_custom_op/flash_attention/attention.py +0 -406
- mindspore/ops/_op_impl/_custom_op/flash_attention/constants.py +0 -41
- mindspore/ops/_op_impl/_custom_op/flash_attention/flash_attention_bwd.py +0 -467
- mindspore/ops/_op_impl/_custom_op/flash_attention/flash_attention_fwd.py +0 -563
- mindspore/ops/_op_impl/_custom_op/flash_attention/flash_attention_impl.py +0 -193
- mindspore/ops/_op_impl/_custom_op/flash_attention/tik_ops_utils.py +0 -435
- mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/__init__.py +0 -0
- mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/sparse_tiling.py +0 -45
- mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/strategy.py +0 -67
- mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/wukong_tiling.py +0 -62
- {mindspore-2.2.10.dist-info → mindspore-2.2.14.dist-info}/WHEEL +0 -0
- {mindspore-2.2.10.dist-info → mindspore-2.2.14.dist-info}/entry_points.txt +0 -0
- {mindspore-2.2.10.dist-info → mindspore-2.2.14.dist-info}/top_level.txt +0 -0
|
@@ -227,7 +227,11 @@ class ExpandDims(PrimitiveWithCheck):
|
|
|
227
227
|
def infer_value(self, input_x, axis):
|
|
228
228
|
value = None
|
|
229
229
|
if input_x is not None and axis is not None:
|
|
230
|
-
|
|
230
|
+
dtype = input_x.dtype
|
|
231
|
+
if input_x.dtype == mstype.bfloat16:
|
|
232
|
+
cpu_cast = Cast().set_device("CPU")
|
|
233
|
+
input_x = cpu_cast(input_x, mstype.float32)
|
|
234
|
+
value = Tensor(np.expand_dims(input_x.asnumpy(), axis), dtype)
|
|
231
235
|
return value
|
|
232
236
|
|
|
233
237
|
|
|
@@ -375,6 +379,9 @@ class Cast(PrimitiveWithCheck):
|
|
|
375
379
|
if isinstance(x, (int, float)):
|
|
376
380
|
value = Tensor(np.array(x).astype(np_dst_type), dtype=dst_type)
|
|
377
381
|
else:
|
|
382
|
+
if x.dtype == mstype.bfloat16:
|
|
383
|
+
cpu_cast = Cast().set_device("CPU")
|
|
384
|
+
x = cpu_cast(x, mstype.float32)
|
|
378
385
|
value = Tensor(x.asnumpy().astype(np_dst_type), dtype=dst_type)
|
|
379
386
|
return value
|
|
380
387
|
|
|
@@ -5837,6 +5844,9 @@ class SpaceToBatchND(Primitive):
|
|
|
5837
5844
|
``Ascend`` ``GPU`` ``CPU``
|
|
5838
5845
|
|
|
5839
5846
|
Examples:
|
|
5847
|
+
>>> import mindspore
|
|
5848
|
+
>>> from mindspore import Tensor, ops
|
|
5849
|
+
>>> import numpy as np
|
|
5840
5850
|
>>> block_shape = [2, 2]
|
|
5841
5851
|
>>> paddings = [[0, 0], [0, 0]]
|
|
5842
5852
|
>>> space_to_batch_nd = ops.SpaceToBatchND(block_shape, paddings)
|
|
@@ -6279,8 +6289,8 @@ class Sort(Primitive):
|
|
|
6279
6289
|
Sorts the elements of the input tensor along the given dimension in the specified order.
|
|
6280
6290
|
|
|
6281
6291
|
.. warning::
|
|
6282
|
-
Currently, the data types of
|
|
6283
|
-
|
|
6292
|
+
Currently, the data types of float16, uint8, int8, int16, int32, int64 are well supported.
|
|
6293
|
+
If use float32, it may cause loss of accuracy.
|
|
6284
6294
|
|
|
6285
6295
|
Args:
|
|
6286
6296
|
axis (int, optional): The dimension to sort along. Default: ``-1``, means the last dimension.
|
|
@@ -8719,7 +8729,8 @@ class TopK(Primitive):
|
|
|
8719
8729
|
- GPU: float16, float32.
|
|
8720
8730
|
- CPU: all numeric types.
|
|
8721
8731
|
|
|
8722
|
-
- **k** (int) - The number of top elements to be computed along the last dimension
|
|
8732
|
+
- **k** (Union(Tensor, int)) - The number of top elements to be computed along the last dimension.
|
|
8733
|
+
If `k` is a Tensor, the supported dtype is int32 and it should be 0-D or 1-D with shape :math:`(1, )` .
|
|
8723
8734
|
|
|
8724
8735
|
Outputs:
|
|
8725
8736
|
A tuple consisting of `values` and `indexes`.
|
|
@@ -470,7 +470,7 @@ class Custom(ops.PrimitiveWithInfer):
|
|
|
470
470
|
op_path_in_cache = [] # Save paths for op functions created in the cached.
|
|
471
471
|
custom_aot_warning = True # Flag to enable warnings about custom aot path white list
|
|
472
472
|
|
|
473
|
-
def __init__(self, func, out_shape=None, out_dtype=None, func_type=
|
|
473
|
+
def __init__(self, func, out_shape=None, out_dtype=None, func_type="hybrid", bprop=None, reg_info=None):
|
|
474
474
|
super().__init__("Custom")
|
|
475
475
|
|
|
476
476
|
self.supported_targets = [ASCEND, GPU, CPU]
|
|
@@ -472,7 +472,7 @@ class Print(Primitive):
|
|
|
472
472
|
|
|
473
473
|
Examples:
|
|
474
474
|
>>> import numpy as np
|
|
475
|
-
>>> from mindspore import Tensor, nn
|
|
475
|
+
>>> from mindspore import Tensor, nn, ops
|
|
476
476
|
>>> class PrintDemo(nn.Cell):
|
|
477
477
|
... def __init__(self):
|
|
478
478
|
... super(PrintDemo, self).__init__()
|
|
@@ -388,7 +388,7 @@ class NonMaxSuppressionV3(Primitive):
|
|
|
388
388
|
single score associated with each box (i.e., each row of the `boxes` Tensor).
|
|
389
389
|
It is required that the number of scores in `scores` must be equal to the number of boxes in `boxes`.
|
|
390
390
|
The supported data type is float32.
|
|
391
|
-
- **max_output_size** (Union[Tensor, Number.
|
|
391
|
+
- **max_output_size** (Union[Tensor, Number.int]) - A scalar integer Tensor representing the maximum
|
|
392
392
|
number of boxes to be selected by non max suppression. The supported data type is int32.
|
|
393
393
|
- **iou_threshold** (Union[Tensor, Number.Float]) - A scalar float Tensor represents the threshold
|
|
394
394
|
used for determining if the intersection over union (IOU) between boxes is too high.
|
|
@@ -459,7 +459,7 @@ class NonMaxSuppressionWithOverlaps(Primitive):
|
|
|
459
459
|
single score associated with each box (i.e., each row of the `boxes` Tensor).
|
|
460
460
|
It is required that the number of scores in `scores` must be equal to the number of boxes in `boxes`.
|
|
461
461
|
The supported data type is float32.
|
|
462
|
-
- **max_output_size** (Union[Tensor, Number.
|
|
462
|
+
- **max_output_size** (Union[Tensor, Number.int]) - A scalar integer Tensor representing the maximum
|
|
463
463
|
number of boxes to be selected by non max suppression, and max_output_size must be equal to or greater
|
|
464
464
|
than 0.
|
|
465
465
|
Types allowed:int32.
|
|
@@ -816,7 +816,7 @@ class ResizeBicubic(Primitive):
|
|
|
816
816
|
Examples:
|
|
817
817
|
>>> import mindspore
|
|
818
818
|
>>> import numpy as np
|
|
819
|
-
>>> from mindspore import Tensor, ops
|
|
819
|
+
>>> from mindspore import Tensor, ops, nn
|
|
820
820
|
>>> class NetResizeBicubic(nn.Cell):
|
|
821
821
|
... def __init__(self):
|
|
822
822
|
... super(NetResizeBicubic, self).__init__()
|
|
@@ -642,6 +642,55 @@ class FusedAdaFactorWithGlobalNorm(FusedAdaFactor):
|
|
|
642
642
|
return param_type
|
|
643
643
|
|
|
644
644
|
|
|
645
|
+
class GenerateEodMask(Primitive):
|
|
646
|
+
r"""
|
|
647
|
+
Given the input `inputs_ids`, if found eod_token_id, the output position and attention mask matrix will be reset.
|
|
648
|
+
This means the `position_id` will start counting from 0, and the corresponding mask matrix will be filled with 0.
|
|
649
|
+
|
|
650
|
+
Args:
|
|
651
|
+
eod_token_id (int) - In the NLP scenario, this value corresponds to the id of
|
|
652
|
+
the symbol of 'EodOfDocument' in the vocabulary.
|
|
653
|
+
|
|
654
|
+
Inputs:
|
|
655
|
+
- **inputs_ids** (Tensor) - token id, a 2-D Tensor with shape :math:`(batch\_size, seq\_length)`.
|
|
656
|
+
|
|
657
|
+
Outputs:
|
|
658
|
+
- **position_id** (Tensor) - position id matrix with same shape and type as original `inputs_ids`.
|
|
659
|
+
- **attention_mask** (Tensor) - attention mask matrix with type
|
|
660
|
+
float16 and shape :math:`(batch\_size, seq\_length)`.
|
|
661
|
+
|
|
662
|
+
Supported Platforms:
|
|
663
|
+
``Ascend``
|
|
664
|
+
|
|
665
|
+
Examples:
|
|
666
|
+
>>> op = ops.GenerateEodMask(eod_token_id=0)
|
|
667
|
+
>>> position, mask = op(Tensor([[1, 0, 3], [1, 0, 0]], dtype=mindspore.int32))
|
|
668
|
+
>>> print(position)
|
|
669
|
+
[[0 1 0] [0 0 1]]
|
|
670
|
+
>>> print(mask)
|
|
671
|
+
[[[ 1. 0. 0.]
|
|
672
|
+
[1. 1. 0.]
|
|
673
|
+
[0. 0. 1.]]
|
|
674
|
+
[[1. 0. 0.]
|
|
675
|
+
[0. 1. 0.]
|
|
676
|
+
[0. 1. 1.]]]
|
|
677
|
+
|
|
678
|
+
Raises:
|
|
679
|
+
- **TypeError** - If `eod_token_id` is not int.
|
|
680
|
+
- **TypeError** - If `inputs_ids` is not int.
|
|
681
|
+
- **ValueError** - If `inputs_ids` is not a 2-D Tensor.
|
|
682
|
+
"""
|
|
683
|
+
@prim_attr_register
|
|
684
|
+
def __init__(self, n_pos, eod_token_id, n_step, n_error_mode='specific'):
|
|
685
|
+
"""Initialize GenerateEodMask"""
|
|
686
|
+
validator.check_value_type("eod_token_id", eod_token_id, [int], self.name)
|
|
687
|
+
validator.check_value_type("n_pos", n_pos, [int], self.name)
|
|
688
|
+
validator.check_value_type("n_step", n_step, [list], self.name)
|
|
689
|
+
validator.check_value_type("n_error_mode", n_error_mode, [str], self.name)
|
|
690
|
+
self.init_prim_io_names(inputs=['inputs_ids'],
|
|
691
|
+
outputs=['position_ids'])
|
|
692
|
+
|
|
693
|
+
|
|
645
694
|
class ScaleGrad(PrimitiveWithInfer):
|
|
646
695
|
"""
|
|
647
696
|
Scale the input grad according to the loss scale.
|
|
@@ -123,6 +123,64 @@ class _MathBinaryOp(_BinaryOp):
|
|
|
123
123
|
real_shape = [dim if cmp_dim > 0 else cmp_dim for dim, cmp_dim in zip(shape_value, cmp_shape)]
|
|
124
124
|
return tuple(real_shape)
|
|
125
125
|
|
|
126
|
+
class SilentCheck(Primitive):
|
|
127
|
+
"""
|
|
128
|
+
Implement SilentCheck on `pre_val`, `min_val`, `max_val`, `result` and
|
|
129
|
+
update them inplace with given parameters.
|
|
130
|
+
|
|
131
|
+
Args:
|
|
132
|
+
c_min_steps (int): an int determines...
|
|
133
|
+
|
|
134
|
+
c_thresh_l1 (float): a float determines...
|
|
135
|
+
|
|
136
|
+
c_coeff_l1 (float): a float determines...
|
|
137
|
+
|
|
138
|
+
c_thresh_l2 (float): a float determines...
|
|
139
|
+
|
|
140
|
+
c_coeff_l2 (float): a float determines...
|
|
141
|
+
|
|
142
|
+
Inputs:
|
|
143
|
+
- **val** (Tensor) - Tensor with dtype float32.
|
|
144
|
+
- **input_grad** (Parameter) - Tensor with dtype float32.
|
|
145
|
+
- **pre_val** (Parameter) - Input Parameter with dtype float32.
|
|
146
|
+
- **min_val** (Parameter) - Input Parameter with dtype float32.
|
|
147
|
+
- **max_val** (Parameter) - Input Parameter with dtype float32.
|
|
148
|
+
- **val_counter** (Parameter) - Input Parameter with dtype int32.
|
|
149
|
+
|
|
150
|
+
Outputs:
|
|
151
|
+
Tuple of 5 Tensors, the updated parameters.
|
|
152
|
+
- **input_grad** (Tensor) - Tensor with dtype float32.
|
|
153
|
+
- **pre_val** (Tensor) - Tensor with dtype float32.
|
|
154
|
+
- **min_val** (Tensor) - Tensor with dtype float32.
|
|
155
|
+
- **max_val** (Tensor) - Tensor with dtype float32.
|
|
156
|
+
- **result** (Tensor) - Tensor with dtype int32.
|
|
157
|
+
|
|
158
|
+
Raises:
|
|
159
|
+
TypeError: If `val` is not Tensor with dtype float32.
|
|
160
|
+
TypeError: If `result` is not Tensor with dtype int32.
|
|
161
|
+
TypeError: If `pre_val`, `min_val`, `max_val`, `input_grad` are not all Parameter type with dtype float32.
|
|
162
|
+
TypeError: If `c_thresh_l1` or `c_coeff_l1` is not a float number.
|
|
163
|
+
TypeError: If `c_min_steps` is not an int number.
|
|
164
|
+
|
|
165
|
+
Supported Platforms:
|
|
166
|
+
``Ascend``
|
|
167
|
+
|
|
168
|
+
Examples:
|
|
169
|
+
>>> from mindspore.ops.operations.math_ops import SilentCheck
|
|
170
|
+
>>> silent_check = SilentCheck()
|
|
171
|
+
xxx
|
|
172
|
+
"""
|
|
173
|
+
|
|
174
|
+
@prim_attr_register
|
|
175
|
+
def __init__(self, c_min_steps, c_thresh_l1, c_coeff_l1, c_thresh_l2, c_coeff_l2):
|
|
176
|
+
"""Initialize SilentCheck."""
|
|
177
|
+
validator.check_value_type("c_min_steps", c_min_steps, [int], self.name)
|
|
178
|
+
validator.check_value_type("c_thresh_l1", c_thresh_l1, [float], self.name)
|
|
179
|
+
validator.check_value_type("c_coeff_l1", c_coeff_l1, [float], self.name)
|
|
180
|
+
validator.check_value_type("c_thresh_l2", c_thresh_l2, [float], self.name)
|
|
181
|
+
validator.check_value_type("c_coeff_l2", c_coeff_l2, [float], self.name)
|
|
182
|
+
self.add_prim_attr('side_effect_mem', True)
|
|
183
|
+
|
|
126
184
|
|
|
127
185
|
class _BitwiseBinaryOp(_MathBinaryOp):
|
|
128
186
|
"""
|
|
@@ -462,6 +520,7 @@ class AssignAdd(Primitive):
|
|
|
462
520
|
>>> import mindspore
|
|
463
521
|
>>> import numpy as np
|
|
464
522
|
>>> from mindspore import Tensor, ops, nn
|
|
523
|
+
>>> from mindspore.common.initializer import initializer
|
|
465
524
|
>>> class Net(nn.Cell):
|
|
466
525
|
... def __init__(self):
|
|
467
526
|
... super(Net, self).__init__()
|
|
@@ -512,6 +571,7 @@ class AssignSub(Primitive):
|
|
|
512
571
|
>>> import mindspore
|
|
513
572
|
>>> import numpy as np
|
|
514
573
|
>>> from mindspore import Tensor, ops, nn
|
|
574
|
+
>>> from mindspore.common.initializer import initializer
|
|
515
575
|
>>> class Net(nn.Cell):
|
|
516
576
|
... def __init__(self):
|
|
517
577
|
... super(Net, self).__init__()
|
|
@@ -6569,9 +6629,9 @@ class LinSpace(Primitive):
|
|
|
6569
6629
|
|
|
6570
6630
|
Inputs:
|
|
6571
6631
|
- **start** (Tensor) - Start value of interval, 0-D Tensor with dtype float32 or float64.
|
|
6572
|
-
- **stop** (Tensor) - Last value of interval, 0-D Tensor with dtype
|
|
6573
|
-
- **num** (int) - Number of ticks in the interval, inclusive of `start` and `stop`.
|
|
6574
|
-
|
|
6632
|
+
- **stop** (Tensor) - Last value of interval, 0-D Tensor with dtype float32 or float64.
|
|
6633
|
+
- **num** (Union[int, Tensor]) - Number of ticks in the interval, inclusive of `start` and `stop`.
|
|
6634
|
+
Must be a positive integer. When the input is Tensor, it must be a 0-D Tensor with dtype int32 or int64.
|
|
6575
6635
|
|
|
6576
6636
|
Outputs:
|
|
6577
6637
|
Tensor, has the same shape and dtype as `start`.
|
|
@@ -7253,6 +7313,7 @@ class Igamma(Primitive):
|
|
|
7253
7313
|
|
|
7254
7314
|
Examples:
|
|
7255
7315
|
>>> import numpy as np
|
|
7316
|
+
>>> import mindspore
|
|
7256
7317
|
>>> from mindspore import Tensor, ops
|
|
7257
7318
|
>>> a = Tensor(np.array([2.0, 4.0, 6.0, 8.0]).astype(np.float32))
|
|
7258
7319
|
>>> x = Tensor(np.array([2.0, 3.0, 4.0, 5.0]).astype(np.float32))
|
|
@@ -7291,6 +7352,7 @@ class Igammac(Primitive):
|
|
|
7291
7352
|
``Ascend`` ``GPU`` ``CPU``
|
|
7292
7353
|
|
|
7293
7354
|
Examples:
|
|
7355
|
+
>>> import mindspore
|
|
7294
7356
|
>>> import numpy as np
|
|
7295
7357
|
>>> from mindspore import Tensor, ops
|
|
7296
7358
|
>>> a = Tensor(np.array([2.0, 4.0, 6.0, 8.0]).astype(np.float32))
|
|
@@ -3777,7 +3777,7 @@ class LayerNorm(Primitive):
|
|
|
3777
3777
|
- **output_x** (Tensor) - The normalized input, has the same type and shape as the `input_x`.
|
|
3778
3778
|
- **mean** (Tensor) - The first `begin_norm_axis` dimensions of `mean` shape is the same as `input_x`,
|
|
3779
3779
|
and the remaining dimensions are 1. Suppose the shape of the `input_x` is :math:`(x_1, x_2, \ldots, x_R)`,
|
|
3780
|
-
the shape of the `mean` is :math:`(x_1, \ldots, x_{
|
|
3780
|
+
the shape of the `mean` is :math:`(x_1, \ldots, x_{begin\_params\_axis}, 1, \ldots, 1)`
|
|
3781
3781
|
(when `begin_params_axis=0`, the shape of `mean` is :math:`(1, \ldots, 1)` ).
|
|
3782
3782
|
- **variance** (Tensor) - Shape is the same as `mean` .
|
|
3783
3783
|
|
|
@@ -4917,6 +4917,7 @@ class Adam(Primitive):
|
|
|
4917
4917
|
>>> import mindspore
|
|
4918
4918
|
>>> import numpy as np
|
|
4919
4919
|
>>> from mindspore import Tensor, nn, ops
|
|
4920
|
+
>>> from mindspore import Parameter
|
|
4920
4921
|
>>> class Net(nn.Cell):
|
|
4921
4922
|
... def __init__(self):
|
|
4922
4923
|
... super(Net, self).__init__()
|
|
@@ -9991,6 +9992,9 @@ class FractionalMaxPool3DWithFixedKsize(Primitive):
|
|
|
9991
9992
|
``Ascend`` ``GPU`` ``CPU``
|
|
9992
9993
|
|
|
9993
9994
|
Examples:
|
|
9995
|
+
>>> import numpy as np
|
|
9996
|
+
>>> from mindspore import Tensor, ops
|
|
9997
|
+
>>> from mindspore import dtype as mstype
|
|
9994
9998
|
>>> x = Tensor(np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16])
|
|
9995
9999
|
... .reshape([1, 1, 2, 2, 4]), mstype.float32)
|
|
9996
10000
|
>>> random_samples = Tensor(np.array([0.7, 0.7, 0.7]).reshape([1, 1, 3]), mstype.float32)
|
|
@@ -11363,7 +11367,7 @@ class PromptFlashAttention(Primitive):
|
|
|
11363
11367
|
For each element, 0 indicates retention and 1 indicates discard. Input tensor of shape :math:`(B, 1, S, S)`.
|
|
11364
11368
|
- **actual_seq_lengths** (Tensor): Describe actual sequence length of each input with data type of int.
|
|
11365
11369
|
- **actual_seq_lengths_kv** (Tensor): Describe actual sequence length of each input with data type of int.
|
|
11366
|
-
- **
|
|
11370
|
+
- **pse_shift** (Tensor) - The position encoding tensor with data type of float16 or float32.
|
|
11367
11371
|
- **dep_scale1** (Tensor)
|
|
11368
11372
|
- **quant_scale1** (Tensor)
|
|
11369
11373
|
- **deq_scale2** (Tensor)
|
|
@@ -11406,7 +11410,7 @@ class PromptFlashAttention(Primitive):
|
|
|
11406
11410
|
validator.check_value_type('num_key_value_heads', num_key_value_heads, [int], self.name)
|
|
11407
11411
|
validator.check_value_type('sparse_mode', sparse_mode, [int], self.name)
|
|
11408
11412
|
self.init_prim_io_names(inputs=["query", "key", "value", "attn_mask", "actual_seq_lengths",
|
|
11409
|
-
"actual_seq_lengths_kv", "
|
|
11413
|
+
"actual_seq_lengths_kv", "pse_shift", "deq_scale1", "quant_scale1",
|
|
11410
11414
|
"deq_scale2", "quant_scale2", "quant_offset2"],
|
|
11411
11415
|
outputs=["attention_out"])
|
|
11412
11416
|
|
|
@@ -11417,41 +11421,50 @@ class FlashAttentionScore(Primitive):
|
|
|
11417
11421
|
.. warning::
|
|
11418
11422
|
This is an experimental API that is subject to change or deletion.
|
|
11419
11423
|
B -- Batch size
|
|
11420
|
-
|
|
11421
|
-
|
|
11422
|
-
|
|
11423
|
-
|
|
11424
|
+
S1 -- Sequence length of query
|
|
11425
|
+
S2 -- Sequence length of key and value
|
|
11426
|
+
N1 -- Num heads of query
|
|
11427
|
+
N2 -- Num heads of key and value, and N2 must be a factor of N1
|
|
11428
|
+
D -- head size
|
|
11429
|
+
H1 -- Hidden size of query, which equals to N1 * D
|
|
11430
|
+
H2 -- Hidden size of key and value, which equals to N2 * D
|
|
11424
11431
|
Args:
|
|
11425
|
-
head_num (int): The
|
|
11432
|
+
head_num (int): The head num of query.
|
|
11426
11433
|
keep_prob (float): The keep probability of dropout. Default: 1.0.
|
|
11427
11434
|
scale_value (float): The scale value. Default: 1.0.
|
|
11428
11435
|
pre_tokens (int): Previous tokens. Default: 65536.
|
|
11429
11436
|
next_tokens (int): Next tokens. Default: 65536.
|
|
11430
11437
|
inner_precise (int): Specify the execution mode, where 0 indicates high precision mode and 1 indicates high
|
|
11431
|
-
performance mode. Default: 0.
|
|
11438
|
+
performance mode. Only support 0 currently. Default: 0.
|
|
11432
11439
|
input_layout (str, optional): Specifies the layout of `query`, the value must be one of ["BSH", "BNSD"].
|
|
11433
11440
|
Default: "BSH".
|
|
11434
11441
|
sparse_mode (int): Default 0.
|
|
11435
11442
|
|
|
11436
11443
|
Inputs:
|
|
11437
|
-
- **query** (Tensor) - The query tensor
|
|
11438
|
-
Input tensor of shape :math:`(B,
|
|
11439
|
-
- **key** (Tensor) - The key tensor
|
|
11440
|
-
Input tensor of shape :math:`(B,
|
|
11441
|
-
- **value** (Tensor) - The value tensor
|
|
11442
|
-
Input tensor of shape :math:`(B,
|
|
11443
|
-
- **
|
|
11444
|
-
|
|
11445
|
-
- **drop_mask** (Tensor) - The dropout mask tensor
|
|
11446
|
-
Input tensor of shape :math:`(B,
|
|
11447
|
-
- **real_shift** (None) - The position embedding code of float16 or float32, not implemented yet.
|
|
11444
|
+
- **query** (Tensor[float16, float32, bfloat16]) - The query tensor.
|
|
11445
|
+
Input tensor of shape :math:`(B, S1, H1)` or `(B, N1, S1, D)`.
|
|
11446
|
+
- **key** (Tensor[float16, float32, bfloat16]) - The key tensor.
|
|
11447
|
+
Input tensor of shape :math:`(B, S2, H2)` or `(B, N2, S2, D)`.
|
|
11448
|
+
- **value** (Tensor[float16, float32, bfloat16]) - The value tensor.
|
|
11449
|
+
Input tensor of shape :math:`(B, S2, H2)` or `(B, N2, S2, D)`.
|
|
11450
|
+
- **real_shift** (Tensor[float16, float32, bfloat16], None) - The position embedding code.
|
|
11451
|
+
Input tensor of shape :math: `(B, N1, S1, S2)` or `(B, N1, 1, S2)`.
|
|
11452
|
+
- **drop_mask** (Tensor[uint8], None) - The dropout mask tensor.
|
|
11453
|
+
Input tensor of shape :math:`(B, N1, S1, S2 // 8) or None`.
|
|
11448
11454
|
- **padding_mask** (None) - The padding mask of float16 or float32, not implemented yet.
|
|
11449
|
-
- **
|
|
11455
|
+
- **attn_mask** (Tensor[uint8], None) - The attention mask tensor.
|
|
11456
|
+
For each element, 0 indicates retention and 1 indicates discard.
|
|
11457
|
+
Input tensor of shape :math:`(B, N1, S1, S2)`, `(B, 1, S1, S2)` or `(S1, S2)`.
|
|
11458
|
+
- **prefix** (Tensor[int64], None) - Not implemented yet.
|
|
11459
|
+
Input tensor of shape :math:`(B,)`.
|
|
11450
11460
|
|
|
11451
11461
|
Outputs:
|
|
11452
|
-
- **
|
|
11453
|
-
- **
|
|
11454
|
-
- **
|
|
11462
|
+
- **softmax_max** (Tensor[float32]) - (B, N1, S1, 8)
|
|
11463
|
+
- **softmax_sum** (Tensor[float32]) - (B, N1, S1, 8)
|
|
11464
|
+
- **softmax_out** (Tensor[float32]) - Useless output, ignore it. Output tensor of shape : `()`
|
|
11465
|
+
- **attention_out** (Tensor[float16, float32, bfloat16]) - The output of attention, its shape, and data type
|
|
11466
|
+
are the same as the query.
|
|
11467
|
+
|
|
11455
11468
|
Supported Platforms:
|
|
11456
11469
|
``Ascend``
|
|
11457
11470
|
"""
|
|
@@ -11469,14 +11482,14 @@ class FlashAttentionScore(Primitive):
|
|
|
11469
11482
|
validator.check_value_type('next_tokens', next_tokens, [int], self.name)
|
|
11470
11483
|
validator.check_value_type('inner_precise', inner_precise, [int], self.name)
|
|
11471
11484
|
validator.check_value_type('sparse_mode', sparse_mode, [int], self.name)
|
|
11472
|
-
if inner_precise not in [0
|
|
11473
|
-
raise ValueError(f"Attribute 'inner_precise' must be
|
|
11485
|
+
if inner_precise not in [0]:
|
|
11486
|
+
raise ValueError(f"Attribute 'inner_precise' must be 0, but got {inner_precise}")
|
|
11474
11487
|
validator.check_value_type('input_layout', input_layout, [str], self.name)
|
|
11475
11488
|
if input_layout not in ["BSH", "BNSD"]:
|
|
11476
11489
|
raise ValueError(f"Attribute 'input_layout' must be either 'BSH' or 'BNSD', but got {input_layout}")
|
|
11477
11490
|
self.init_prim_io_names(
|
|
11478
|
-
inputs=['query', 'key', 'value', '
|
|
11479
|
-
outputs=['
|
|
11491
|
+
inputs=['query', 'key', 'value', 'real_shift', 'drop_mask', 'padding_mask', 'attn_mask', 'prefix'],
|
|
11492
|
+
outputs=['softmax_max', 'softmax_sum', 'softmax_out', 'attention_out'])
|
|
11480
11493
|
|
|
11481
11494
|
|
|
11482
11495
|
class RmsNorm(Primitive):
|
|
@@ -11514,3 +11527,57 @@ class RmsNorm(Primitive):
|
|
|
11514
11527
|
"""Initialize Dense."""
|
|
11515
11528
|
validator.check_value_type("epsilon", epsilon, [float], self.name)
|
|
11516
11529
|
self.init_prim_io_names(inputs=['x', 'gamma'], outputs=["y", "rstd"])
|
|
11530
|
+
|
|
11531
|
+
|
|
11532
|
+
class PagedAttention(Primitive):
|
|
11533
|
+
r"""
|
|
11534
|
+
.. warning::
|
|
11535
|
+
This is an experimental API that is subject to change or deletion.
|
|
11536
|
+
"""
|
|
11537
|
+
@prim_attr_register
|
|
11538
|
+
def __init__(self, head_num, scale_value=1.0, kv_head_num=0):
|
|
11539
|
+
"""Initialize PagedAttention"""
|
|
11540
|
+
validator.check_value_type('head_num', head_num, [int], self.name)
|
|
11541
|
+
validator.check_value_type('scale_value', scale_value, [float], self.name) # scale after qkbmm
|
|
11542
|
+
validator.check_value_type('kv_head_num', kv_head_num, [int], self.name) # for MQA
|
|
11543
|
+
self.init_prim_io_names(
|
|
11544
|
+
inputs=['query', 'key_cache', 'value_cache', 'block_tables', 'context_lens'],
|
|
11545
|
+
outputs=['attention_out'])
|
|
11546
|
+
|
|
11547
|
+
|
|
11548
|
+
class PagedAttentionMask(Primitive):
|
|
11549
|
+
r"""
|
|
11550
|
+
.. warning::
|
|
11551
|
+
This is an experimental API that is subject to change or deletion.
|
|
11552
|
+
"""
|
|
11553
|
+
@prim_attr_register
|
|
11554
|
+
def __init__(self, head_num, scale_value=1.0, kv_head_num=0):
|
|
11555
|
+
"""Initialize PagedAttentionMask"""
|
|
11556
|
+
validator.check_value_type('head_num', head_num, [int], self.name)
|
|
11557
|
+
validator.check_value_type('scale_value', scale_value, [float], self.name) # scale after qkbmm
|
|
11558
|
+
validator.check_value_type('kv_head_num', kv_head_num, [int], self.name) # for MQA
|
|
11559
|
+
self.init_prim_io_names(
|
|
11560
|
+
inputs=['query', 'key_cache', 'value_cache', 'block_tables', 'context_lens', 'alibi_mask'],
|
|
11561
|
+
outputs=['attention_out'])
|
|
11562
|
+
|
|
11563
|
+
|
|
11564
|
+
class ReshapeAndCache(Primitive):
|
|
11565
|
+
r"""
|
|
11566
|
+
.. warning::
|
|
11567
|
+
This is an experimental API that is subject to change or deletion.
|
|
11568
|
+
"""
|
|
11569
|
+
__mindspore_signature__ = (
|
|
11570
|
+
sig.make_sig('key', dtype=sig.sig_dtype.T),
|
|
11571
|
+
sig.make_sig('value', dtype=sig.sig_dtype.T),
|
|
11572
|
+
sig.make_sig('key_cache', sig.sig_rw.RW_WRITE, dtype=sig.sig_dtype.T),
|
|
11573
|
+
sig.make_sig('value_cache', sig.sig_rw.RW_WRITE, dtype=sig.sig_dtype.T),
|
|
11574
|
+
sig.make_sig('slot_mapping', dtype=sig.sig_dtype.T1),
|
|
11575
|
+
)
|
|
11576
|
+
|
|
11577
|
+
@prim_attr_register
|
|
11578
|
+
def __init__(self):
|
|
11579
|
+
"""Initialize ReshapeAndCache"""
|
|
11580
|
+
self.init_prim_io_names(
|
|
11581
|
+
inputs=['key', 'value', 'key_cache', 'value_cache', 'slot_mapping'],
|
|
11582
|
+
outputs=['key_out'])
|
|
11583
|
+
self.add_prim_attr('side_effect_mem', True)
|
|
@@ -479,8 +479,8 @@ class SparseToDenseV2(Primitive):
|
|
|
479
479
|
Tensor, converted from sparse tensor. The dtype is same as `values`, and the shape is `output_shape`.
|
|
480
480
|
|
|
481
481
|
Raises:
|
|
482
|
-
TypeError: If the dtype of `indices` is neither
|
|
483
|
-
TypeError: If the dtype of `outputshape` is neither
|
|
482
|
+
TypeError: If the dtype of `indices` is neither int32 nor int64.
|
|
483
|
+
TypeError: If the dtype of `outputshape` is neither int32 nor int64.
|
|
484
484
|
ValueError: If the shape of `output_shape`, shape of `indices`,
|
|
485
485
|
shape of `default_value` and shape of `values` don't meet the parameter description.
|
|
486
486
|
ValueError: If each Element of `output_shape` is not > 0.
|
|
@@ -2382,8 +2382,8 @@ class SparseCountSparseOutput(Primitive):
|
|
|
2382
2382
|
Args:
|
|
2383
2383
|
binary_output (bool) - If ``False`` , output the number of occurrences of each value,
|
|
2384
2384
|
if ``True`` output 1 for orresponding values. Default: ``False`` .
|
|
2385
|
-
minlength(Scalar) -
|
|
2386
|
-
maxlength(Scalar) -
|
|
2385
|
+
minlength(Scalar) - int type minimum value to count, Default: ``-1`` .
|
|
2386
|
+
maxlength(Scalar) - int type maximum value to count, Default: ``-1`` .
|
|
2387
2387
|
|
|
2388
2388
|
Inputs:
|
|
2389
2389
|
- **indices** (Tensor) - Tensor representing the position of the element in the sparse
|
|
@@ -0,0 +1,162 @@
|
|
|
1
|
+
# Copyright 2024 Huawei Technologies Co., Ltd
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
# ============================================================================
|
|
15
|
+
"""Silent Check."""
|
|
16
|
+
import os
|
|
17
|
+
|
|
18
|
+
from mindspore.common.tensor import Tensor
|
|
19
|
+
from mindspore.common.parameter import Parameter
|
|
20
|
+
import mindspore.common.dtype as mstype
|
|
21
|
+
|
|
22
|
+
from . import operations
|
|
23
|
+
from .operations._inner_ops import _MirrorSilentCheck
|
|
24
|
+
from .operations import RmsNorm as OriginRmsNorm
|
|
25
|
+
from .operations import LayerNorm as OriginLayerNorm
|
|
26
|
+
from .primitive import Primitive
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
NPU_ASD_ENABLE = 'NPU_ASD_ENABLE'
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class ASDBase:
|
|
33
|
+
"""
|
|
34
|
+
ASDBase is the base class of operator with accuracy-sensitive detection feature in python.
|
|
35
|
+
|
|
36
|
+
Args:
|
|
37
|
+
cls (Primitive): Original operator requiring accuracy-sensitive detection feature.
|
|
38
|
+
args (tuple): A variable parameter tuple to the original operator.
|
|
39
|
+
kwargs (dict): A variable parameter dictionary passed the original operator.
|
|
40
|
+
|
|
41
|
+
Supported Platforms:
|
|
42
|
+
``Ascend``
|
|
43
|
+
|
|
44
|
+
Examples:
|
|
45
|
+
>>> from mindspore.ops.silent_check import ASDBase
|
|
46
|
+
>>> from mindspore.ops import LayerNorm as OriginLayerNorm
|
|
47
|
+
>>> class LayerNormASD(ASDBase):
|
|
48
|
+
... def __init__(self, *args, **kwargs):
|
|
49
|
+
... super().__init__(OriginLayerNorm, *args, **kwargs)
|
|
50
|
+
... # init parameters for accuracy-sensitive detection by calling the base class method generate_params()
|
|
51
|
+
... self.pre_val, self.min_val, self.max_val, self.cnt = self.generate_params()
|
|
52
|
+
...
|
|
53
|
+
... def __call__(self, input_x, gamma, beta):
|
|
54
|
+
... if self.enable_check:
|
|
55
|
+
... # execute accuracy-sensitive detection by calling the check_op of base class
|
|
56
|
+
... input_x = self.check_op(
|
|
57
|
+
... input_x, self.pre_val, self.min_val, self.max_val, self.cnt, None)
|
|
58
|
+
... self.cnt += 1
|
|
59
|
+
... # return the result of original operator
|
|
60
|
+
... return self.op(input_x, gamma, beta)
|
|
61
|
+
"""
|
|
62
|
+
_index = 0
|
|
63
|
+
__ms_class__ = True
|
|
64
|
+
|
|
65
|
+
def __init__(self, cls, *args, **kwargs):
|
|
66
|
+
self.op = cls(*args, **kwargs)
|
|
67
|
+
self.check_op = _MirrorSilentCheck()
|
|
68
|
+
self._suffix = "ASD_" + cls.__name__
|
|
69
|
+
primitive_attr = dir(Primitive)
|
|
70
|
+
self._op_attr_dict = {
|
|
71
|
+
name for name in primitive_attr if not name.startswith("_")}
|
|
72
|
+
self.enable_check = os.environ.get(NPU_ASD_ENABLE) == "1"
|
|
73
|
+
|
|
74
|
+
def __getattr__(self, name):
|
|
75
|
+
def method_wrapper(*args, **kwargs):
|
|
76
|
+
out = getattr(self.op, name)(*args, **kwargs)
|
|
77
|
+
if out is self.op:
|
|
78
|
+
return self
|
|
79
|
+
return out
|
|
80
|
+
|
|
81
|
+
if name in self._op_attr_dict:
|
|
82
|
+
if callable(getattr(self.op, name)):
|
|
83
|
+
return method_wrapper
|
|
84
|
+
if hasattr(self.op, name):
|
|
85
|
+
return getattr(self.op, name)
|
|
86
|
+
return super().__getattr__(self, name)
|
|
87
|
+
|
|
88
|
+
def __repr__(self):
|
|
89
|
+
return self.op.__repr__()
|
|
90
|
+
|
|
91
|
+
def generate_params(self):
|
|
92
|
+
"""
|
|
93
|
+
Generate support params for accuracy-sensitive detection.
|
|
94
|
+
|
|
95
|
+
Returns:
|
|
96
|
+
tuple consisting of four elements.
|
|
97
|
+
The derived class initializes the parameters required for accuracy-sensitive detection by calling
|
|
98
|
+
this function.
|
|
99
|
+
|
|
100
|
+
Examples:
|
|
101
|
+
>>> from mindspore.ops.silent_check import ASDBase
|
|
102
|
+
>>> from mindspore.ops import LayerNorm as OriginLayerNorm
|
|
103
|
+
>>> class LayerNormASD(ASDBase):
|
|
104
|
+
... def __init__(self, *args, **kwargs):
|
|
105
|
+
... super().__init__(OriginLayerNorm, *args, **kwargs)
|
|
106
|
+
... # init parameters for accuracy-sensitive detection by calling the base class function
|
|
107
|
+
... self.pre_val, self.min_val, self.max_val, self.cnt = self.generate_params()
|
|
108
|
+
"""
|
|
109
|
+
pre_val = Parameter(Tensor(0, mstype.float32),
|
|
110
|
+
name=f"{self._suffix}_pre_val_{self._index}",
|
|
111
|
+
requires_grad=False)
|
|
112
|
+
min_val = Parameter(Tensor(0, mstype.float32),
|
|
113
|
+
name=f"{self._suffix}_min_val_{self._index}",
|
|
114
|
+
requires_grad=False)
|
|
115
|
+
max_val = Parameter(Tensor(0, mstype.float32),
|
|
116
|
+
name=f"{self._suffix}_max_val_{self._index}",
|
|
117
|
+
requires_grad=False)
|
|
118
|
+
cnt = Parameter(Tensor(0, mstype.int32),
|
|
119
|
+
name=f"{self._suffix}_cnt_{self._index}",
|
|
120
|
+
requires_grad=False)
|
|
121
|
+
ASDBase._index += 1
|
|
122
|
+
return pre_val, min_val, max_val, cnt
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
class RmsNormASD(ASDBase):
|
|
126
|
+
"""
|
|
127
|
+
RmsNorm with ASD.
|
|
128
|
+
"""
|
|
129
|
+
|
|
130
|
+
def __init__(self, *args, **kwargs):
|
|
131
|
+
super().__init__(OriginRmsNorm, *args, **kwargs)
|
|
132
|
+
self.pre_val, self.min_val, self.max_val, self.cnt = self.generate_params()
|
|
133
|
+
|
|
134
|
+
def __call__(self, input_x, gamma):
|
|
135
|
+
if self.enable_check:
|
|
136
|
+
input_x = self.check_op(
|
|
137
|
+
input_x, self.pre_val, self.min_val, self.max_val, self.cnt, None)
|
|
138
|
+
self.cnt += 1
|
|
139
|
+
return self.op(input_x, gamma)
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
class LayerNormASD(ASDBase):
|
|
143
|
+
"""
|
|
144
|
+
LayerNorm with ASD.
|
|
145
|
+
"""
|
|
146
|
+
|
|
147
|
+
def __init__(self, *args, **kwargs):
|
|
148
|
+
super().__init__(OriginLayerNorm, *args, **kwargs)
|
|
149
|
+
self.pre_val, self.min_val, self.max_val, self.cnt = self.generate_params()
|
|
150
|
+
|
|
151
|
+
def __call__(self, input_x, gamma, beta):
|
|
152
|
+
if self.enable_check:
|
|
153
|
+
input_x = self.check_op(
|
|
154
|
+
input_x, self.pre_val, self.min_val, self.max_val, self.cnt, None)
|
|
155
|
+
self.cnt += 1
|
|
156
|
+
return self.op(input_x, gamma, beta)
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
def _silent_check():
|
|
160
|
+
if os.environ.get(NPU_ASD_ENABLE) == "1":
|
|
161
|
+
operations.LayerNorm = LayerNormASD
|
|
162
|
+
operations.RmsNorm = RmsNormASD
|