diff --git a/.github/workflows/ci-platform-gap9.yml b/.github/workflows/ci-platform-gap9.yml index 280812d9..b5994416 100644 --- a/.github/workflows/ci-platform-gap9.yml +++ b/.github/workflows/ci-platform-gap9.yml @@ -64,13 +64,35 @@ jobs: docker-image: ${{ needs.select-env.outputs.image }} pytest-marker: "training" + # On-chip training: weights and activations stay in L2, nothing spills to L3. + # Without this job the L2 training entries are configured but never executed. + gap9-training-tiled-l2-singlebuffer: + needs: select-env + uses: ./.github/workflows/_runner-gap9-tiled.yml + with: + runner: ${{ needs.select-env.outputs.runner }} + docker-image: ${{ needs.select-env.outputs.image }} + pytest-markers: "gap9_tiled and training and singlebuffer and l2" + gap9-training-tiled-l3-singlebuffer: needs: select-env uses: ./.github/workflows/_runner-gap9-tiled.yml with: runner: ${{ needs.select-env.outputs.runner }} docker-image: ${{ needs.select-env.outputs.image }} - pytest-markers: "gap9_tiled and training and singlebuffer and l3" + pytest-markers: "gap9_tiled and training and singlebuffer and l3 and not recompute" + + # Gradient-checkpointed training: the same models and budgets as the singlebuffer + # job, but replaying a solved recompute schedule so checkpointed activations are + # regenerated before the backward pass rather than held live. Its own job so it + # does not stretch the singlebuffer job's runtime and a failure is attributable. + gap9-training-tiled-l3-recompute: + needs: select-env + uses: ./.github/workflows/_runner-gap9-tiled.yml + with: + runner: ${{ needs.select-env.outputs.runner }} + docker-image: ${{ needs.select-env.outputs.image }} + pytest-markers: "gap9_tiled and training and l3 and recompute" # L3 double-buffered training (CCT) — promote+DB handled by its own job below gap9-training-tiled-l3-doublebuffer: diff --git a/.gitignore b/.gitignore index a26eeb65..bfa200ab 100644 --- a/.gitignore +++ b/.gitignore @@ -58,6 +58,9 @@ DeeployTest/Tests/**/pactIntegerizationProto DeeployTest/Tests/**/quantlib DeeployTest/Tests/**/*.py DeeployTest/Tests/**/*.json +# Recompute schedules are inputs, not generated artefacts: a test that names one +# and cannot find it fails at generation, which is how it reached CI unnoticed. +!DeeployTest/Tests/**/recompute_*.json DeeployTest/Tests/**/generateTest.py DeeployTest/out.txt diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a8f323fa..2877364c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -65,6 +65,9 @@ repos: - id: clang-format name: Autoformat C/C++ Files args: ["-i"] + # The hook's default file set includes .json, and reformatting a recompute + # schedule leaves it unparseable. Only C and C++ here. + types_or: [c, c++] stages: [pre-commit, pre-merge-commit, pre-push, manual] - repo: https://github.com/adrienverge/yamllint.git rev: v1.33.0 diff --git a/Deeploy/DeeployTypes.py b/Deeploy/DeeployTypes.py index 035111d4..d1a76834 100644 --- a/Deeploy/DeeployTypes.py +++ b/Deeploy/DeeployTypes.py @@ -979,8 +979,12 @@ def hoistConstant(self, Returns the name of the newly registed ConstantBuffer """ - assert len(constant.outputs) <= 1, f"Constant {constant.name} has more than one output" - + # A ConstantBuffer is global and read-only, so several nodes reading one is + # sound and needs no copy. Requiring a single consumer blocked every graph + # where a frozen weight is read twice: a quantised weight reaching both the + # forward MatMul and its backward Gemm, and a transposed weight shared the + # same way after constant folding. Duplicating to satisfy the assertion is + # not free either, since it charges the full constant again per consumer. name = name if name is not None else constant.name # LMACAN: The shape needs to be copied into a tuple for pickling to work. Don't ask me why.. @@ -3467,6 +3471,15 @@ def frontEnd(self): log.debug(" - Constant Folding") self._foldConstants(self.graph) + # Folding can hand a single constant to several consumers: a weight reached + # through Constant -> Transpose -> Variable collapses to one Constant, and a + # training graph has the forward MatMul and the backward Gemm both reading it. + # Duplication ran before the fold, so nothing splits that, and hoistConstant + # then asserts. Run it again; it only touches tensors with more than one + # consumer, so it is a no-op when the fold introduced none. + log.debug(" - Duplicate Constants (post-fold)") + self._duplicateConstants(self.graph) + log.info(f"> Export State to {_middlewarePreLoweringFilename}[.onnx|.pkl]") self.exportDeeployState(self.deeployStateDir, _middlewarePreLoweringFilename) diff --git a/Deeploy/Targets/GAP9/Bindings.py b/Deeploy/Targets/GAP9/Bindings.py index b273687e..60079e50 100644 --- a/Deeploy/Targets/GAP9/Bindings.py +++ b/Deeploy/Targets/GAP9/Bindings.py @@ -39,13 +39,13 @@ from Deeploy.Targets.PULPOpen.Templates import ConvTemplate, DMASliceTemplate, FloatAddTemplate, \ FloatAveragePoolTemplate, FloatBatchNormTemplate, FloatConvGradTemplate, FloatConvTemplate, FloatGELUTemplate, \ FloatGemmTemplate, FloatGlobalAveragePoolTemplate, FloatInPlaceAccumulatorV2Template, FloatLayernormTemplate, \ - FloatMatMulTemplate, FloatMaxPoolTemplate, FloatMulTemplate, FloatReluTemplate, FloatSoftmaxTemplate, \ - GEMMTemplate, MatrixVectorTemplate, MaxPoolTemplate, MSELossTemplate, MulTemplate, ReduceMeanTemplate, \ - RequantShiftTemplate, ReshapeTemplate, RQAddTemplate, RQSiHardswishTemplate, SGDTemplate, \ + FloatMatMulDequantTemplate, FloatMatMulTemplate, FloatMaxPoolTemplate, FloatMulTemplate, FloatReluTemplate, \ + FloatSoftmaxTemplate, GEMMTemplate, MatrixVectorTemplate, MaxPoolTemplate, MSELossTemplate, MulTemplate, \ + ReduceMeanTemplate, RequantShiftTemplate, ReshapeTemplate, RQAddTemplate, RQSiHardswishTemplate, SGDTemplate, \ SoftmaxCrossEntropyLossTemplate, TallGEMMTemplate, TransposeTemplate, UniformRequantShiftTemplate, \ iRMSNormTemplate, iSoftmaxTemplate -from Deeploy.Targets.PULPOpen.TypeCheckers import PULPConvChecker, PULPLinearChecker, PULPMaxPoolChecker, \ - PULPRequantShiftChecker +from Deeploy.Targets.PULPOpen.TypeCheckers import PULPConvChecker, PULPDequantConvChecker, PULPDequantGEMMChecker, \ + PULPDequantMatMulChecker, PULPLinearChecker, PULPMaxPoolChecker, PULPRequantShiftChecker from Deeploy.TilingExtension.CodeTransformationPasses.TilingVariableReplacement import TilingVariableReplacement, \ TilingVariableReplacementUpdate @@ -157,6 +157,12 @@ AddTemplate.referenceTemplate, GAP9Transformer) for type1 in IntegerDataTypes for type2 in IntegerDataTypes +] + [ + # fp32 activations against an int8 weight whose Dequant has been folded in, so + # the dequantised weight is never materialised. GAP9 keeps its own binding list; + # adding this to PULPMatMulBindings alone has no effect here. + NodeBinding(PULPDequantMatMulChecker([PointerClass(float32_t), PointerClass(int8_t)], [PointerClass(float32_t)]), + FloatMatMulDequantTemplate.referenceTemplate, GAP9Transformer) ] + [ NodeBinding(AddChecker([PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]), FloatAddTemplate.referenceTemplate, GAP9Transformer) @@ -196,6 +202,22 @@ ] GAP9FloatGEMMBindings = [ + # int8 weight with the Dequant folded in. Needed on the BACKWARD Gemm too, not + # just the forward MatMul: they share the weight constant, and if the backward + # one selects an fp32 binding, typeInferGlobalCtxt re-types the shared constant + # to float and its allocation quadruples. + NodeBinding( + PULPDequantGEMMChecker( + [PointerClass(float32_t), PointerClass(int8_t), + PointerClass(float32_t)], [PointerClass(float32_t)]), FloatGemmTemplate.referenceDequantTemplate, + GAP9Transformer), + # fp32 first. A binding is matched before a weight constant's type is fixed, and + # typeInferGlobalCtxt then types that constant from the SELECTED binding's + # input_types. With the int8 variant first, an ordinary fp32 weight is claimed by + # it and retyped to int8, silently quantising it: the FP32 Conv kernel test + # failed 512 of 512 and the emitted GAP9 call had no visible declaration. A + # genuinely int8 constant already carries that type, so the fp32 checker rejects + # it and the int8 binding below still wins. NodeBinding( GEMMChecker([PointerClass(float32_t), PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]), FloatGemmTemplate.referenceTemplate, @@ -203,6 +225,14 @@ ] GAP9FloatConv2DBindings = [ + # int8 weight with the Dequant folded in. GAP9 keeps its own binding list, so + # this has to be added here and not only in PULPFloatConv2DBindings. + NodeBinding( + PULPDequantConvChecker( + [PointerClass(float32_t), PointerClass(int8_t), + PointerClass(float32_t)], [PointerClass(float32_t)]), FloatConvTemplate.reference2DIm2ColDequantTemplate, + GAP9Transformer), + # fp32 first, for the reason given on GAP9FloatGEMMBindings above. NodeBinding( ConvChecker([PointerClass(float32_t), PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]), FloatConvTemplate.reference2DIm2ColTemplate, @@ -280,6 +310,13 @@ GAP9MatMulBindings = [ NodeBinding(MatMulChecker([PointerClass(int8_t), PointerClass(int8_t)], [PointerClass(int32_t)]), GEMMTemplate.PULPMM_8_Template, GAP9ClusterTransformer) +] + [ + # int8 weight with the Dequant folded in. PULPDequantMatMulChecker matches only a + # node the fold marked with dequant_scale, so this must come first to be reached + # at all (binding selection stops at the first type check that passes) while + # staying invisible to an ordinary fp32 graph. + NodeBinding(PULPDequantMatMulChecker([PointerClass(float32_t), PointerClass(int8_t)], [PointerClass(float32_t)]), + FloatMatMulDequantTemplate.referenceTemplate, GAP9Transformer) ] + [ NodeBinding(MatMulChecker([PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]), FloatMatMulTemplate.referenceTemplate, GAP9Transformer) diff --git a/Deeploy/Targets/GAP9/Platform.py b/Deeploy/Targets/GAP9/Platform.py index 9e8b89d8..20172513 100644 --- a/Deeploy/Targets/GAP9/Platform.py +++ b/Deeploy/Targets/GAP9/Platform.py @@ -34,12 +34,11 @@ BasicRQIntegerDivBinding from Deeploy.Targets.Generic.Layers import AddLayer, AveragePoolGradLayer, BatchNormalizationGradLayer, \ BatchNormInternalLayer, ConcatLayer, ConvGradBLayer, ConvGradWLayer, ConvGradXLayer, ConvLayer, GatherLayer, \ - GELUGradLayer, GELULayer, GEMMLayer, GlobalAveragePoolGradLayer, GlobalAveragePoolLayer, \ - InPlaceAccumulatorV2Layer, LayerNormGradLayer, LayerNormLayer, MatMulLayer, MaxPoolLayer, MSELossGradLayer, \ - MSELossLayer, MulLayer, PadLayer, QuantLayer, ReduceMeanLayer, ReduceSumLayer, ReluGradLayer, ReluLayer, \ - RequantShiftLayer, ReshapeLayer, RQIntegerDivLayer, RQSiGELULayer, RQSiHardswishLayer, SGDLayer, SliceLayer, \ - SoftmaxCrossEntropyLossGradLayer, SoftmaxCrossEntropyLossLayer, SoftmaxGradLayer, SoftmaxLayer, TransposeLayer, \ - iHardswishLayer, iRMSNormLayer + GELUGradLayer, GELULayer, GlobalAveragePoolGradLayer, GlobalAveragePoolLayer, InPlaceAccumulatorV2Layer, \ + LayerNormGradLayer, LayerNormLayer, MatMulLayer, MaxPoolLayer, MSELossGradLayer, MSELossLayer, MulLayer, PadLayer, \ + QuantLayer, ReduceMeanLayer, ReduceSumLayer, ReluGradLayer, ReluLayer, RequantShiftLayer, ReshapeLayer, \ + RQIntegerDivLayer, RQSiGELULayer, RQSiHardswishLayer, SGDLayer, SliceLayer, SoftmaxCrossEntropyLossGradLayer, \ + SoftmaxCrossEntropyLossLayer, SoftmaxGradLayer, SoftmaxLayer, TransposeLayer, iHardswishLayer, iRMSNormLayer from Deeploy.Targets.Generic.Parsers import AddParser, AveragePool2DParser, BatchNormalizationGradParser, \ BatchNormInternalParser, ConcatParser, Conv2DGradBParser, DequantParser, FlattenParser, GatherParser, \ GELUGradParser, GELUParser, GEMMParser, GlobalAveragePoolGradParser, GlobalAveragePoolParser, \ @@ -52,7 +51,7 @@ from Deeploy.Targets.Generic.Templates import AllocateTemplate as BasicAllocateTemplate from Deeploy.Targets.PULPOpen.Bindings import BasicQuantBindings, PULPDMASliceBindings, PULPDWConv1DBinding, \ PULPReduceMeanBindings, PULPRQSConv1DBindings -from Deeploy.Targets.PULPOpen.Layers import PULPRQSConvLayer, PULPRQSGEMMLayer +from Deeploy.Targets.PULPOpen.Layers import PULPAddLayer, PULPGEMMLayer, PULPRQSConvLayer, PULPRQSGEMMLayer from Deeploy.Targets.PULPOpen.Parsers import PULPConv1DParser, PULPConv2DParser, PULPConvGradW2DParser, \ PULPConvGradX2DParser, PULPDWConv1DParser, PULPDWConv2DParser, PULPDWConvGradW2DParser, PULPDWConvGradX2DParser, \ PULPFPConv2DCHWParser, PULPFPConv2DParser, PULPFPDWConv2DCHWParser, PULPFPDWConv2DParser, PULPGEMMParser, \ @@ -145,7 +144,7 @@ 'RequantizedGemm': PULPRQSGEMMLayer([GAP9_MatrixVecMapper, GAP9_TallGEMMMapper, GAP9_GEMMMapper]), 'Gemm': - GEMMLayer([GAP9_FloatGEMMMapper, GAP9_GEMMDequantMapper]), + PULPGEMMLayer([GAP9_FloatGEMMMapper, GAP9_GEMMDequantMapper]), 'Gelu': GELULayer([GAP9_GELUMapper]), 'GeluGrad': @@ -173,7 +172,7 @@ 'RequantShift': RequantShiftLayer([GAP9_UniformRequantShiftMapper, GAP9_RequantShiftMapper]), 'Add': - AddLayer([GAP9_AddMapper]), + PULPAddLayer([GAP9_AddMapper]), 'Flatten': ReshapeLayer([GAP9_FlattenMapper]), 'Gather': diff --git a/Deeploy/Targets/GAP9/Tiler.py b/Deeploy/Targets/GAP9/Tiler.py index 6c6122d5..ce114947 100644 --- a/Deeploy/Targets/GAP9/Tiler.py +++ b/Deeploy/Targets/GAP9/Tiler.py @@ -39,6 +39,7 @@ from Deeploy.Targets.PULPOpen.TileConstraints.AveragePoolTileConstraint import AveragePoolCTileConstraint from Deeploy.Targets.PULPOpen.TileConstraints.BatchNormTileConstraint import BatchNormalizationGradTileConstraint, \ BatchNormInternalTileConstraint +from Deeploy.Targets.PULPOpen.TileConstraints.BroadcastAddTileConstraint import PULPBroadcastAddTileConstraint from Deeploy.Targets.PULPOpen.TileConstraints.ConvGradConstraint import ConvGradBTileConstraint, \ ConvGradW2DTileConstraint, ConvGradX2DIm2ColHWTileConstraint, DWConvGradW2DTileConstraint, \ DWConvGradX2DTileConstraint, PWConvGradWTileConstraint, PWConvGradXTileConstraint @@ -130,7 +131,7 @@ tileConstraint = TransposeTileConstraint()) GAP9AddTilingReadyBindings = TilingReadyNodeBindings(nodeBindings = GAP9AddBindings, - tileConstraint = AddTileConstraint()) + tileConstraint = PULPBroadcastAddTileConstraint()) GAP9SoftmaxTilingReadyBindings = TilingReadyNodeBindings(nodeBindings = GAP9SoftmaxBindings, tileConstraint = iSoftmaxTileConstraint()) diff --git a/Deeploy/Targets/Generic/Parsers.py b/Deeploy/Targets/Generic/Parsers.py index daa7e397..49ffedd0 100644 --- a/Deeploy/Targets/Generic/Parsers.py +++ b/Deeploy/Targets/Generic/Parsers.py @@ -757,7 +757,13 @@ def parseNodeCtxt(self, self.operatorRepresentation['data_in_1'] = data_in_1.name self.operatorRepresentation['data_in_2'] = data_in_2.name self.operatorRepresentation['data_out'] = data_out.name - self.operatorRepresentation['size'] = np.prod(data_in_1.shape) + # Size the loop by the output, not by an input. With both operands the same + # shape these agree, but a broadcast operand does not: PULPAddLayer leaves a + # Linear's [O] bias at its own shape rather than storing the same row per + # output row, and taking the first input's size then bounded the loop at the + # bias length. The remaining outputs were never written and read back as + # whatever was in memory. + self.operatorRepresentation['size'] = np.prod(data_out.shape) return ctxt, True @@ -2064,6 +2070,13 @@ def parseNode(self, node: gs.Node) -> (bool): self.operatorRepresentation['transB'] = 0 self.operatorRepresentation['transA'] = 0 + # A weight-only-quantised graph can fold its Dequant into this node, in + # which case the B operand arrives as int8 and carries the per-tensor + # affine parameters the kernel needs. Defaults keep the fp32 path + # unchanged for every graph that has no folded Dequant. + self.operatorRepresentation['dequant_scale'] = float(node.attrs.get('dequant_scale', 1.0)) + self.operatorRepresentation['dequant_zero_point'] = int(node.attrs.get('dequant_zero_point', 0)) + return ret def parseNodeCtxt(self, diff --git a/Deeploy/Targets/PULPOpen/Bindings.py b/Deeploy/Targets/PULPOpen/Bindings.py index 69879a41..3e3a8173 100644 --- a/Deeploy/Targets/PULPOpen/Bindings.py +++ b/Deeploy/Targets/PULPOpen/Bindings.py @@ -34,13 +34,13 @@ from Deeploy.Targets.PULPOpen.Templates import ConvTemplate, DMASliceTemplate, FloatAddTemplate, \ FloatAveragePoolTemplate, FloatBatchNormTemplate, FloatConvGradTemplate, FloatConvTemplate, FloatGELUTemplate, \ FloatGemmTemplate, FloatGlobalAveragePoolTemplate, FloatInPlaceAccumulatorV2Template, FloatLayernormTemplate, \ - FloatMatMulTemplate, FloatMaxPoolTemplate, FloatMulTemplate, FloatReduceMeanTemplate, FloatReluTemplate, \ - FloatSoftmaxTemplate, GEMMTemplate, MatrixVectorTemplate, MaxPoolTemplate, MSELossTemplate, MulTemplate, \ - ReduceMeanTemplate, RequantShiftTemplate, ReshapeTemplate, RQAddTemplate, RQSiHardswishTemplate, SGDTemplate, \ - SoftmaxCrossEntropyLossTemplate, TallGEMMTemplate, TransposeTemplate, UniformRequantShiftTemplate, \ + FloatMatMulDequantTemplate, FloatMatMulTemplate, FloatMaxPoolTemplate, FloatMulTemplate, FloatReduceMeanTemplate, \ + FloatReluTemplate, FloatSoftmaxTemplate, GEMMTemplate, MatrixVectorTemplate, MaxPoolTemplate, MSELossTemplate, \ + MulTemplate, ReduceMeanTemplate, RequantShiftTemplate, ReshapeTemplate, RQAddTemplate, RQSiHardswishTemplate, \ + SGDTemplate, SoftmaxCrossEntropyLossTemplate, TallGEMMTemplate, TransposeTemplate, UniformRequantShiftTemplate, \ iRMSNormTemplate, iSoftmaxTemplate -from Deeploy.Targets.PULPOpen.TypeCheckers import PULPConvChecker, PULPLinearChecker, PULPMaxPoolChecker, \ - PULPRequantShiftChecker +from Deeploy.Targets.PULPOpen.TypeCheckers import PULPConvChecker, PULPDequantConvChecker, PULPDequantMatMulChecker, \ + PULPLinearChecker, PULPMaxPoolChecker, PULPRequantShiftChecker from Deeploy.TilingExtension.CodeTransformationPasses.TilingVariableReplacement import TilingVariableReplacement, \ TilingVariableReplacementUpdate @@ -193,6 +193,11 @@ AddTemplate.referenceTemplate, ForkTransformer) for type1 in IntegerDataTypes for type2 in IntegerDataTypes +] + [ + # fp32 activations against an int8 weight: the Dequant has been folded into the + # matmul, so the weight is never materialised as fp32. + NodeBinding(PULPDequantMatMulChecker([PointerClass(float32_t), PointerClass(int8_t)], [PointerClass(float32_t)]), + FloatMatMulDequantTemplate.referenceTemplate, ForkTransformer) ] + [ NodeBinding(AddChecker([PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]), FloatAddTemplate.referenceTemplate, ForkTransformer) @@ -239,6 +244,18 @@ ] PULPFloatConv2DBindings = [ + NodeBinding( + PULPDequantConvChecker( + [PointerClass(float32_t), PointerClass(int8_t), + PointerClass(float32_t)], [PointerClass(float32_t)]), FloatConvTemplate.reference2DIm2ColDequantTemplate, + ForkTransformer), + # fp32 first. A binding is matched before a weight constant's type is fixed, and + # typeInferGlobalCtxt then types that constant from the SELECTED binding's + # input_types. Listing the int8 variant first does not make it "more specific": + # it claims an ordinary fp32 weight and retypes it to int8, silently quantising + # it. That is what failed the FP32 Conv kernel test 512 of 512. A genuinely int8 + # constant already carries that type, so the fp32 checker rejects it and the + # int8 binding below still wins. NodeBinding( ConvChecker([PointerClass(float32_t), PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]), FloatConvTemplate.reference2DIm2ColTemplate, @@ -390,6 +407,17 @@ PULPMatMulBindings = [ NodeBinding(MatMulChecker([PointerClass(int8_t), PointerClass(int8_t)], [PointerClass(int32_t)]), GEMMTemplate.PULPMM_8_Template, ClusterTransformer) +] + [ + # int8 weight with the Dequant folded in. Ordering alone cannot separate this from + # the fp32 binding: selection happens before a weight constant's type is fixed and + # typeInferGlobalCtxt then types that constant from the SELECTED binding, so + # whichever comes first claims the weight and defines its type. int8-first + # silently quantised ordinary fp32 weights (FP32 Conv kernel test, 512 of 512); + # fp32-first inflated folded int8 weights (CCT-QLoRA persistent 409 KB -> 1094 KB). + # PULPDequantMatMulChecker keys on the dequant_scale attribute the fold leaves + # behind instead, so it is invisible to an fp32 graph and can safely come first. + NodeBinding(PULPDequantMatMulChecker([PointerClass(float32_t), PointerClass(int8_t)], [PointerClass(float32_t)]), + FloatMatMulDequantTemplate.referenceTemplate, ForkTransformer) ] + [ NodeBinding(MatMulChecker([PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]), FloatMatMulTemplate.referenceTemplate, ForkTransformer) diff --git a/Deeploy/Targets/PULPOpen/Layers.py b/Deeploy/Targets/PULPOpen/Layers.py index 69ce2fa9..537a0bc5 100644 --- a/Deeploy/Targets/PULPOpen/Layers.py +++ b/Deeploy/Targets/PULPOpen/Layers.py @@ -5,7 +5,7 @@ from typing import List, Tuple from Deeploy.DeeployTypes import NodeMapper, Shape -from Deeploy.Targets.Generic.Layers import RQGEMMLayer, RQSConvLayer +from Deeploy.Targets.Generic.Layers import AddLayer, GEMMLayer, RQGEMMLayer, RQSConvLayer class PULPRQSConvLayer(RQSConvLayer): @@ -41,3 +41,50 @@ def computeShapes(self, inputShapes: Shape, outputShapes: Shape, operatorReprese inputShapes[3] = [inputShapes[1][channelDim]] # Channels out dimension of Kernel return (inputShapes, outputShapes) + + +class PULPGEMMLayer(GEMMLayer): + """GEMM that keeps a one-dimensional bias one-dimensional. + + The generic layer widens the C operand to [M, N] unconditionally. That is what + the kernel used to require, and for a transformer it is expensive: every Linear + bias is O values shared by all M rows, so the widened operand stores the same + vector M times. On CCT-2 at 64 tokens each 128-wide bias occupies 32 KB to carry + 512 B, and the eight of them live at the peak account for 252 KB, around 13% of + the measured peak under LoRA and linear probing. + + PULPFloatGEMMTemplate now passes the bias row stride to the kernel, so a + one-dimensional C is read with stride 0 and needs no widening. A C operand that + arrives with a real second dimension is left alone and behaves as before. + """ + + def __init__(self, maps: List[NodeMapper]): + super().__init__(maps) + + def computeShapes(self, inputShapes: Shape, outputShapes: Shape, operatorRepresentation, + channels_first) -> Tuple[Shape, Shape]: + if len(inputShapes) == 3 and len([d for d in inputShapes[2] if d != 1]) <= 1: + return (inputShapes, outputShapes) # broadcast bias: leave as [O] + return super().computeShapes(inputShapes, outputShapes, operatorRepresentation, channels_first) + + +class PULPAddLayer(AddLayer): + """Add that leaves a broadcast second operand at its own shape. + + The generic layer rewrites the lower-rank operand to the higher-rank one's + shape, which turns a Linear's [O] bias into a full [1, M, O] tensor storing the + same row M times. PULPBroadcastAddTileConstraint and PULPFloatAddTemplate + handle the unequal shapes, so the rewrite is only needed when the operand is + not a broadcast. + """ + + def __init__(self, maps: List[NodeMapper]): + super().__init__(maps) + + def computeShapes(self, inputShapes: Shape, outputShapes: Shape, operatorRepresentation, + channels_first) -> Tuple[Shape, Shape]: + big, small = (0, 1) if len(inputShapes[0]) >= len(inputShapes[1]) else (1, 0) + if len(inputShapes[big]) != len(inputShapes[small]) and \ + len([d for d in inputShapes[small] if d != 1]) <= 1: + return (inputShapes, [inputShapes[big]]) + return super().computeShapes(inputShapes, outputShapes, operatorRepresentation, channels_first) diff --git a/Deeploy/Targets/PULPOpen/Platform.py b/Deeploy/Targets/PULPOpen/Platform.py index 0e59976b..55199ded 100644 --- a/Deeploy/Targets/PULPOpen/Platform.py +++ b/Deeploy/Targets/PULPOpen/Platform.py @@ -15,7 +15,7 @@ BasicRQIntegerDivBinding from Deeploy.Targets.Generic.Layers import AddLayer, AveragePoolGradLayer, AveragePoolLayer, \ BatchNormalizationGradLayer, BatchNormInternalLayer, ConcatLayer, ConvGradBLayer, ConvGradWLayer, ConvGradXLayer, \ - ConvLayer, GatherLayer, GELUGradLayer, GELULayer, GEMMLayer, GlobalAveragePoolGradLayer, GlobalAveragePoolLayer, \ + ConvLayer, GatherLayer, GELUGradLayer, GELULayer, GlobalAveragePoolGradLayer, GlobalAveragePoolLayer, \ InPlaceAccumulatorV2Layer, LayerNormGradLayer, LayerNormLayer, MatMulLayer, MaxPoolGradLayer, MaxPoolLayer, \ MSELossGradLayer, MSELossLayer, MulLayer, PadLayer, QuantLayer, ReduceMeanLayer, ReduceSumLayer, ReluGradLayer, \ ReluLayer, RequantShiftLayer, ReshapeLayer, RQIntegerDivLayer, RQSiGELULayer, RQSiHardswishLayer, SGDLayer, \ @@ -35,7 +35,7 @@ MergeConstAddAndRequantPass, MergeTrueIntegerDivRequantShiftPass, QuantPatternPass, RQSSplitPass, \ SkipEmptyConcatPass, SkipUnityRequantPass, iGELURequantMergePass, iHardswishRequantMergePass from Deeploy.Targets.PULPOpen.Bindings import BasicQuantBindings, PULPDMASliceBindings, PULPDWConv1DBinding -from Deeploy.Targets.PULPOpen.Layers import PULPRQSConvLayer, PULPRQSGEMMLayer +from Deeploy.Targets.PULPOpen.Layers import PULPAddLayer, PULPGEMMLayer, PULPRQSConvLayer, PULPRQSGEMMLayer from Deeploy.Targets.PULPOpen.Parsers import PULPConv1DParser, PULPConv2DParser, PULPConvGradW2DParser, \ PULPConvGradX2DParser, PULPDWConv1DParser, PULPDWConv2DParser, PULPDWConvGradW2DParser, PULPDWConvGradX2DParser, \ PULPFPConv2DParser, PULPFPDWConv2DParser, PULPGEMMParser, PULPMatrixVecParser, PULPPWConvGradW2DParser, \ @@ -62,8 +62,9 @@ PULPSoftmaxCrossEntropyGradTilingReadyBindings, PULPSoftmaxCrossEntropyTilingReadyBindings, \ PULPSoftmaxGradTilingReadyBindings, PULPSoftmaxTilingReadyBindings, PULPTransposeTilingReadyBindings, \ PULPUniformRQSTilingReadyBindings -from Deeploy.Targets.PULPOpen.TopologyOptimizationPasses.Passes import PULPAddRequantMergePass, \ - PULPConvRequantMergePass, PULPGEMMRequantMergePass, PULPMatMulRequantMergePass, TransposeGemmSquashPass +from Deeploy.Targets.PULPOpen.TopologyOptimizationPasses.Passes import FoldDequantIntoMatMulPass, \ + PULPAddRequantMergePass, PULPConvRequantMergePass, PULPGEMMRequantMergePass, PULPMatMulRequantMergePass, \ + TransposeGemmSquashPass from Deeploy.Targets.PULPOpen.TopologyOptimizationPasses.SplitConvGradPass import SplitConvGradPass RQAddMapper = NodeMapper(RQAddParser(), PULPRQAddTilingReadyBindings) @@ -151,7 +152,7 @@ 'ConvGradB': ConvGradBLayer([ConvGradBMapper]), 'RequantizedConv': PULPRQSConvLayer([Conv2DMapper, DWConv2DMapper, Conv1DMapper, DWConv1DMapper]), 'RequantizedGemm': PULPRQSGEMMLayer([MatrixVecMapper, TallGEMMMapper, GEMMMapper]), - 'Gemm': GEMMLayer([FloatGEMMMapper, GEMMDequantMapper]), + 'Gemm': PULPGEMMLayer([FloatGEMMMapper, GEMMDequantMapper]), 'Gelu': GELULayer([GELUMapper]), 'GeluGrad': GELUGradLayer([GELUGradMapper]), 'LayerNormalization': LayerNormLayer([LayerNormMapper]), @@ -175,7 +176,7 @@ 'ReduceMean': ReduceMeanLayer([ReduceMeanMapper]), 'ReduceSum': ReduceSumLayer([ReduceSumMapper]), 'RequantShift': RequantShiftLayer([UniformRequantShiftMapper, RequantShiftMapper]), - 'Add': AddLayer([AddMapper]), + 'Add': PULPAddLayer([AddMapper]), 'Flatten': ReshapeLayer([FlattenMapper]), 'Gather': GatherLayer([GatherMapper]), 'Mul': MulLayer([MulMapper]), @@ -271,29 +272,34 @@ class PULPStructBuffer(StructBuffer): deallocTemplate = NodeTemplate("") -PULPOptimizer = TopologyOptimizer([ - SplitConvGradPass(), - TransposeGemmSquashPass(), - QuantPatternPass(), - DequantPatternPass(), - SkipEmptyConcatPass(), - SkipUnityRequantPass(previous_op_regex = "Concat", num_inputs = 2), - SkipUnityRequantPass(previous_op_regex = "Reshape|Transpose", num_inputs = 1), - SkipUnityRequantPass(previous_op_regex = "Reshape|Transpose", num_inputs = 1), - RQSSplitPass(), - MergeTrueIntegerDivRequantShiftPass(), - IntegerDivRequantMergePass(), - iGELURequantMergePass(), - iHardswishRequantMergePass(), - PULPConvRequantMergePass(), - MergeConstAddAndRequantPass(), - PULPGEMMRequantMergePass(), - PULPMatMulRequantMergePass(), - PULPAddRequantMergePass(), - RemoveEmptyConvBiasPass(), - RemoveOnlySingletonReduceMeanPass(), -], - name = "PULPOptimizer") +PULPOptimizer = TopologyOptimizer( + [ + SplitConvGradPass(), + TransposeGemmSquashPass(), + QuantPatternPass(), + DequantPatternPass(), + # Fold a weight's Dequant into the matmul that reads it, so the dequantised + # weight is never materialised. Placed straight after DequantPatternPass, which + # is what normalises the Dequant nodes this matches on. + FoldDequantIntoMatMulPass(), + SkipEmptyConcatPass(), + SkipUnityRequantPass(previous_op_regex = "Concat", num_inputs = 2), + SkipUnityRequantPass(previous_op_regex = "Reshape|Transpose", num_inputs = 1), + SkipUnityRequantPass(previous_op_regex = "Reshape|Transpose", num_inputs = 1), + RQSSplitPass(), + MergeTrueIntegerDivRequantShiftPass(), + IntegerDivRequantMergePass(), + iGELURequantMergePass(), + iHardswishRequantMergePass(), + PULPConvRequantMergePass(), + MergeConstAddAndRequantPass(), + PULPGEMMRequantMergePass(), + PULPMatMulRequantMergePass(), + PULPAddRequantMergePass(), + RemoveEmptyConvBiasPass(), + RemoveOnlySingletonReduceMeanPass(), + ], + name = "PULPOptimizer") # SCHEREMO: stdint is included before pulp_nn_kernels.h because it is supposed to be included in there, but isn't... _includeList = [ diff --git a/Deeploy/Targets/PULPOpen/Templates/FloatAddTemplate.py b/Deeploy/Targets/PULPOpen/Templates/FloatAddTemplate.py index 200ad1b9..65707791 100644 --- a/Deeploy/Targets/PULPOpen/Templates/FloatAddTemplate.py +++ b/Deeploy/Targets/PULPOpen/Templates/FloatAddTemplate.py @@ -2,9 +2,46 @@ # # SPDX-License-Identifier: Apache-2.0 -from Deeploy.DeeployTypes import NodeTemplate +from typing import Dict, List, Tuple -referenceTemplate = NodeTemplate(""" +from Deeploy.DeeployTypes import NetworkContext, NodeTemplate, OperatorRepresentation + + +def _isBroadcast(shape) -> bool: + return len([dim for dim in shape if dim != 1]) <= 1 + + +class PULPFloatAddTemplate(NodeTemplate): + + def alignToContext(self, ctxt: NetworkContext, + operatorRepresentation: OperatorRepresentation) -> Tuple[NetworkContext, Dict, List[str]]: + # An operand with a single non-unit axis is a vector added to every row of + # the other. Keeping it that way avoids storing the same row once per output + # row; the loop below reads it with a wrapped index instead. + # + # Either operand can be the broadcast one. A Linear exported as MatMul + Add + # puts the bias second, but the same graph after lowering can present it + # first, and indexing a two-element bias by the output index reads past it. + # Normalise here so the template always adds data_in_2 into data_in_1. + shape1 = ctxt.lookup(operatorRepresentation['data_in_1']).shape + shape2 = ctxt.lookup(operatorRepresentation['data_in_2']).shape + + if _isBroadcast(shape2) and not _isBroadcast(shape1): + broadcastShape = shape2 + elif _isBroadcast(shape1) and not _isBroadcast(shape2): + operatorRepresentation['data_in_1'], operatorRepresentation['data_in_2'] = \ + operatorRepresentation['data_in_2'], operatorRepresentation['data_in_1'] + broadcastShape = shape1 + else: + broadcastShape = None + + operatorRepresentation['broadcast'] = broadcastShape is not None + if 'rowLen' not in operatorRepresentation: + operatorRepresentation['rowLen'] = (broadcastShape[-1] if broadcastShape else (shape2[-1] if shape2 else 1)) + return ctxt, operatorRepresentation, [] + + +referenceTemplate = PULPFloatAddTemplate(""" // Add Parallel with 1x6 unrolling (Name: ${nodeName}, Op: ${nodeOp}) uint8_t ${nodeName}_core_id = (uint8_t) pi_core_id(); uint8_t ${nodeName}_log2Core = (uint8_t) log2(NUM_CORES); @@ -12,6 +49,18 @@ uint32_t ${nodeName}_chunk_start = (uint32_t) MIN(${nodeName}_chunk*${nodeName}_core_id, (uint32_t) ${size}); uint32_t ${nodeName}_chunk_stop = (uint32_t) MIN(${nodeName}_chunk_start + ${nodeName}_chunk, (uint32_t) ${size}); +% if broadcast: +// data_in_2 holds one row of ${rowLen} shared by every output row. Walk the row +// index alongside i rather than dividing: there is no hardware divide here. +uint32_t ${nodeName}_j = ${nodeName}_chunk_start % ${rowLen}; +for (uint32_t i = ${nodeName}_chunk_start; i < ${nodeName}_chunk_stop; i++) { + ${data_out}[i] = ${data_in_1}[i] + ${data_in_2}[${nodeName}_j]; + ${nodeName}_j++; + if (${nodeName}_j == ${rowLen}) { + ${nodeName}_j = 0; + } +} +% else: uint32_t i = ${nodeName}_chunk_start; for (; i + 5 < ${nodeName}_chunk_stop; i += 6) { ${data_out}[i] = ${data_in_1}[i] + ${data_in_2}[i]; @@ -25,4 +74,5 @@ for (; i < ${nodeName}_chunk_stop; i++) { ${data_out}[i] = ${data_in_1}[i] + ${data_in_2}[i]; } -""") \ No newline at end of file +% endif +""") diff --git a/Deeploy/Targets/PULPOpen/Templates/FloatConvTemplate.py b/Deeploy/Targets/PULPOpen/Templates/FloatConvTemplate.py index cf5e2165..cfa900e3 100644 --- a/Deeploy/Targets/PULPOpen/Templates/FloatConvTemplate.py +++ b/Deeploy/Targets/PULPOpen/Templates/FloatConvTemplate.py @@ -255,3 +255,40 @@ def hoistTransientBuffers(self, ctxt: NetworkContext, ref_${data_out}_${data_out} += ${ch_im_out} * ${dim_im_out_x} * ${dim_im_out_y}; } """) + +# 2D conv whose weight arrives int8, with the Dequant folded in. Identical body to +# reference2DIm2ColTemplate, bound to the kernel that dequantises inside the loop +# and given the per-tensor scale and zero point. +reference2DIm2ColDequantTemplate = PULP2DFloatConvIm2ColTemplate(""" +// 2D FP Conv HWC with Im2Col and ChannelOout parallelism (Name: ${nodeName}, Op: ${nodeOp}) + +${data_in_type.typeName} ref_${data_out}_${data_in} = ${data_in}; +${data_out_type.typeName} ref_${data_out}_${data_out} = ${data_out}; + +for (uint32_t n=0; n<${batch}; ++n) { + PULP_Conv2d_Im2Col_fp32_i8_fp32_HWC( + ref_${data_out}_${data_in}, + ${dim_im_in_x}, + ${dim_im_in_y}, + ${ch_im_in}, + ${weight}, + ${ch_im_out}, + ${dim_kernel_x}, + ${dim_kernel_y}, + ${stride_x}, + ${stride_y}, + ${bias}, ${has_bias}, + ref_${data_out}_${data_out}, + ${padding_y_top}, + ${padding_y_bottom}, + ${padding_x_left}, + ${padding_x_right}, + ${ctxtBuffer}, + ${context.get('dequant_scale', 1.0)}f, + ${context.get('dequant_zero_point', 0)} + ); + + ref_${data_out}_${data_in} += ${ch_im_in} * ${dim_im_in_x} * ${dim_im_in_y}; + ref_${data_out}_${data_out} += ${ch_im_out} * ${dim_im_out_x} * ${dim_im_out_y}; +} +""") diff --git a/Deeploy/Targets/PULPOpen/Templates/FloatGemmTemplate.py b/Deeploy/Targets/PULPOpen/Templates/FloatGemmTemplate.py index ef046f19..0f7bddd3 100644 --- a/Deeploy/Targets/PULPOpen/Templates/FloatGemmTemplate.py +++ b/Deeploy/Targets/PULPOpen/Templates/FloatGemmTemplate.py @@ -22,6 +22,18 @@ def alignToContext(self, ctxt: NetworkContext, operatorRepresentation['C'] = None operatorRepresentation['C_type'] = PointerClass(float32_t) # Default to fp32 type operatorRepresentation['C_batched'] = False + operatorRepresentation['C_stride'] = 0 + return ctxt, operatorRepresentation, [] + + # Row stride of the bias. A bias the parser left one-dimensional holds O + # values shared by every output row, so the kernel re-reads the same vector + # (stride 0) instead of the graph carrying M identical copies of it. A + # two-dimensional bias keeps the original [M,O] walk. + biasShape = ctxt.lookup(operatorRepresentation['C']).shape + broadcast = len([d for d in biasShape if d != 1]) <= 1 + operatorRepresentation['C_stride'] = 0 if broadcast else operatorRepresentation['O'] + if broadcast: + operatorRepresentation['C_batched'] = False return ctxt, operatorRepresentation, [] @@ -48,10 +60,66 @@ def alignToContext(self, ctxt: NetworkContext, ${N}, ${O}, ${transA}, - ${transB} + ${transB}, + ${C_stride} ); % else: PULP_Gemm_fp${A_type.referencedType.typeWidth}_fp${B_type.referencedType.typeWidth}_fp${C_type.referencedType.typeWidth}_fp${data_out_type.referencedType.typeWidth}( + ref_${data_out}_${A}, + ref_${data_out}_${B}, + NULL, + ref_${data_out}_${data_out}, + ${M}, + ${N}, + ${O}, + ${transA}, + ${transB}, + ${C_stride} + ); + % endif + % if A_batched: + ref_${data_out}_${A} += ${M} * ${N}; + % endif + + % if B_batched: + ref_${data_out}_${B} += ${N} * ${O}; + % endif + + % if C is not None and C_batched: + ref_${data_out}_${C} += ${M} * ${O}; + % endif + + ref_${data_out}_${data_out} += ${M} * ${O}; +} +""") + +# Gemm whose weight operand is int8 with the Dequant folded in. +referenceDequantTemplate = PULPFloatGEMMTemplate(""" +// GEMM (Name: ${nodeName}, Op: ${nodeOp}) +${A_type.typeName} ref_${data_out}_${A} = ${A}; +${B_type.typeName} ref_${data_out}_${B} = ${B}; +% if C is not None: +${C_type.typeName} ref_${data_out}_${C} = ${C}; +% else: +${C_type.typeName} ref_${data_out}_C = NULL; +% endif +${data_out_type.typeName} ref_${data_out}_${data_out} = ${data_out}; + +for(uint32_t i=0; i<${batch}; i++){ + % if C is not None: + PULP_Gemm_fp32_i8_fp32_fp32( + ref_${data_out}_${A}, + ref_${data_out}_${B}, + ref_${data_out}_${C}, + ref_${data_out}_${data_out}, + ${M}, + ${N}, + ${O}, + ${transA}, + ${transB} + ); + % else: + PULP_Gemm_fp32_i8_fp32_fp32( ref_${data_out}_${A}, ref_${data_out}_${B}, NULL, @@ -77,4 +145,4 @@ def alignToContext(self, ctxt: NetworkContext, ref_${data_out}_${data_out} += ${M} * ${O}; } -""") \ No newline at end of file +""") diff --git a/Deeploy/Targets/PULPOpen/Templates/FloatMatMulDequantTemplate.py b/Deeploy/Targets/PULPOpen/Templates/FloatMatMulDequantTemplate.py new file mode 100644 index 00000000..093e07ca --- /dev/null +++ b/Deeploy/Targets/PULPOpen/Templates/FloatMatMulDequantTemplate.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: 2025 ETH Zurich and University of Bologna +# +# SPDX-License-Identifier: Apache-2.0 +"""Matmul whose weight operand is int8 and is dequantised inside the kernel. + +A weight-only-quantised model stores the frozen weight as int8 and needs it as +float for the fp32 matmul. Emitting that conversion as a separate Dequant node +materialises the whole dequantised matrix: on CCT-QLoRA, 1088 KB of fp32 tensors +that exist only to be read by the very next node, and that then have to be held +live from the forward pass to the backward one or recomputed. Profiling puts the +cost at 239.95M cycles for the Dequant nodes alone, of which only 48.47M is the +kernel -- the rest is moving those tensors. + +QLoRA does not do this. It dequantises inside the matmul and discards the value, +which is what this template binds: PULP_MatMul_fp32_i8_fp32_unroll1x7 takes the +int8 weight plus the per-tensor scale and zero point and applies them in the loop. +""" + +from Deeploy.DeeployTypes import NodeTemplate + +referenceTemplate = NodeTemplate(""" +// Matmul with an int8 weight, dequantised in-kernel (Name: ${nodeName}, Op: ${nodeOp}) + +for(uint32_t b=0; b<${batch}; b++) { + % if A_batched: + ${A_type.typeName} batch_A = ${A} + b * ${M} * ${N}; + % else: + ${A_type.typeName} batch_A = ${A}; + % endif + + % if B_batched: + ${B_type.typeName} batch_B = ${B} + b * ${N} * ${O}; + % else: + ${B_type.typeName} batch_B = ${B}; + % endif + + ${data_out_type.typeName} batch_out = ${data_out} + b * ${M} * ${O}; + + PULP_MatMul_fp32_i8_fp32_unroll1x7( + batch_A, + batch_B, + batch_out, + ${M}, + ${N}, + ${O}, + ${context.get('dequant_scale', 1.0)}f, + ${context.get('dequant_zero_point', 0)} + ); +} +""") diff --git a/Deeploy/Targets/PULPOpen/TileConstraints/BroadcastAddTileConstraint.py b/Deeploy/Targets/PULPOpen/TileConstraints/BroadcastAddTileConstraint.py new file mode 100644 index 00000000..ad86d2ed --- /dev/null +++ b/Deeploy/Targets/PULPOpen/TileConstraints/BroadcastAddTileConstraint.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: 2025 ETH Zurich and University of Bologna +# +# SPDX-License-Identifier: Apache-2.0 + +from typing import Dict, List, Tuple + +import numpy as np + +from Deeploy.AbstractDataTypes import PointerClass +from Deeploy.CommonExtensions.DataTypes import uint16_t +from Deeploy.DeeployTypes import NetworkContext, OperatorRepresentation +from Deeploy.Targets.Generic.TileConstraints.AddTileConstraint import AddTileConstraint +from Deeploy.TilingExtension.MemoryConstraints import NodeMemoryConstraint +from Deeploy.TilingExtension.TilerModel import TilerModel +from Deeploy.TilingExtension.TilingCodegen import AbsoluteHyperRectangle, HyperRectangle, TilingSchedule, \ + VariableReplacementScheme + + +def isBroadcastOperand(shape) -> bool: + """True when a tensor carries a single non-unit axis, i.e. one row repeated.""" + return len([dim for dim in shape if dim != 1]) <= 1 + + +class PULPBroadcastAddTileConstraint(AddTileConstraint): + """Add where the second operand may be a vector broadcast over the first. + + A Linear exported as MatMul + Add gives the Add a bias of shape [O] against an + activation of shape [1, M, O]. ``AddLayer.computeShapes`` used to rewrite the + bias to the activation's shape, which materialises the same O values M times: + on CCT-2 at 64 tokens that is 32 KB per 128-wide bias holding 512 B of data, + and the biases live at the peak account for 252 KB, around 13% of it. + + Keeping the bias at [O] means the two operands no longer have equal dimensions, + which is the assumption ``BOPTileConstraint`` is built on, in two places: + + * the geometrical constraint ties every axis of both inputs to the output. For + a broadcast operand only the trailing axis has anything to tie; the leading + axes are 1 and stay 1 however the output is tiled. + * the tiling schedule loads the same cube from both inputs. The broadcast + operand instead needs the trailing slice of the output cube, since that is + the part of the vector the tile actually reads. + + Both fall back to the inherited behaviour when the operands do have equal + dimensions, so a genuine elementwise Add is unaffected. + """ + + @classmethod + def addGeometricalConstraint(cls, tilerModel: TilerModel, parseDict: Dict, ctxt: NetworkContext) -> TilerModel: + + inputBuffer2Shape = ctxt.lookup(parseDict[cls.dataIn2Name]).shape + if not isBroadcastOperand(inputBuffer2Shape): + return super().addGeometricalConstraint(tilerModel, parseDict, ctxt) + + inputBuffer1Name = parseDict[cls.dataIn1Name] + inputBuffer2Name = parseDict[cls.dataIn2Name] + outputBufferName = parseDict[cls.dataOutName] + + for bufferName in [inputBuffer1Name, inputBuffer2Name, outputBufferName]: + tilerModel.addTensorDimToModel(ctxt, bufferName) + + input1Shape = ctxt.lookup(inputBuffer1Name).shape + for dim in range(len(input1Shape)): + inputDim1Var = tilerModel.getTensorDimVar(tensorName = inputBuffer1Name, dimIdx = dim) + outputDimVar = tilerModel.getTensorDimVar(tensorName = outputBufferName, dimIdx = dim) + tilerModel.addConstraint(inputDim1Var == outputDimVar) + + # The broadcast operand contributes only its trailing axis, which must match + # the output's trailing axis so a tile reads a whole row of it. + broadcastDimVar = tilerModel.getTensorDimVar(tensorName = inputBuffer2Name, dimIdx = len(inputBuffer2Shape) - 1) + outputLastDimVar = tilerModel.getTensorDimVar(tensorName = outputBufferName, dimIdx = len(input1Shape) - 1) + tilerModel.addConstraint(broadcastDimVar == outputLastDimVar) + + return tilerModel + + @classmethod + def serializeTilingSolution( + cls, tilingSolution: NodeMemoryConstraint, absoluteOutputCubes: List[AbsoluteHyperRectangle], + targetMemLevel: str, ctxt: NetworkContext, + operatorRepresentation: OperatorRepresentation) -> Tuple[VariableReplacementScheme, TilingSchedule]: + + inputBuffer2Shape = ctxt.lookup(operatorRepresentation[cls.dataIn2Name]).shape + if not isBroadcastOperand(inputBuffer2Shape): + return super().serializeTilingSolution(tilingSolution, absoluteOutputCubes, targetMemLevel, ctxt, + operatorRepresentation) + + outputCubes = [cube.rectangle for cube in absoluteOutputCubes] + + addrNames = [cls.dataIn1Name, cls.dataIn2Name, cls.dataOutName] + inputBaseOffsets, outputBaseOffsets = cls.extractBaseAddr(tilingSolution, targetMemLevel, + operatorRepresentation, addrNames) + + replacements = {"size": [], "rowLen": []} + replacementTypes = {"size": PointerClass(uint16_t), "rowLen": PointerClass(uint16_t)} + + inputLoadSchedule = [] + outputLoadSchedule = [] + + for cube in outputCubes: + replacements["size"].append(np.prod(cube.dims)) + replacements["rowLen"].append(cube.dims[-1]) + + # The vector's tile is the trailing slice of the output cube: same + # extent and offset on the last axis, a single element on every other. + rank = len(inputBuffer2Shape) + broadcastOffset = tuple([0] * (rank - 1) + [cube.offset[-1]]) + broadcastDims = tuple([1] * (rank - 1) + [cube.dims[-1]]) + broadcastCube = HyperRectangle(broadcastOffset, broadcastDims) + + inputLoadSchedule.append({cls.dataIn1Name: cube, cls.dataIn2Name: broadcastCube}) + outputLoadSchedule.append({cls.dataOutName: cube}) + + tilingSchedule = TilingSchedule(inputBaseOffsets, outputBaseOffsets, inputLoadSchedule, outputLoadSchedule) + variableReplacementSchedule = VariableReplacementScheme(replacements, replacementTypes) + + return variableReplacementSchedule, tilingSchedule diff --git a/Deeploy/Targets/PULPOpen/TileConstraints/GEMMTileConstraint.py b/Deeploy/Targets/PULPOpen/TileConstraints/GEMMTileConstraint.py index f913b13a..0dc02040 100644 --- a/Deeploy/Targets/PULPOpen/TileConstraints/GEMMTileConstraint.py +++ b/Deeploy/Targets/PULPOpen/TileConstraints/GEMMTileConstraint.py @@ -233,11 +233,21 @@ def addGeometricalConstraint(tilerModel: TilerModel, parseDict: Dict, ctxt: Netw # Add bias constraints only if bias is present if has_bias: - dimOffsetC = len(bufferC.shape) - 2 - addDimVar_1 = tilerModel.getTensorDimVar(tensorName = bufferC.name, dimIdx = dimOffsetC) - addDimVar_2 = tilerModel.getTensorDimVar(tensorName = bufferC.name, dimIdx = dimOffsetC + 1) - tilerModel.addConstraint(outputFirstDimVar == addDimVar_1) - tilerModel.addConstraint(outputSecondDimVar == addDimVar_2) + # A one-dimensional bias is broadcast over the output rows: it has no M + # extent to tie to the output's, and the kernel re-reads the same vector + # for every row (bias row stride 0). Constraining only the trailing + # dimension is what lets it stay [O] instead of being widened to [M,O]. + biasDims = [d for d in bufferC.shape if d != 1] + if len(biasDims) <= 1: + lastDimIdx = len(bufferC.shape) - 1 + addDimVar = tilerModel.getTensorDimVar(tensorName = bufferC.name, dimIdx = lastDimIdx) + tilerModel.addConstraint(outputSecondDimVar == addDimVar) + else: + dimOffsetC = len(bufferC.shape) - 2 + addDimVar_1 = tilerModel.getTensorDimVar(tensorName = bufferC.name, dimIdx = dimOffsetC) + addDimVar_2 = tilerModel.getTensorDimVar(tensorName = bufferC.name, dimIdx = dimOffsetC + 1) + tilerModel.addConstraint(outputFirstDimVar == addDimVar_1) + tilerModel.addConstraint(outputSecondDimVar == addDimVar_2) return tilerModel @@ -340,7 +350,18 @@ def serializeTilingSolution( inputBCubes.append(BCube) if has_bias: - CCube = HyperRectangle(tuple(cube.offset), tuple(cube.dims)) + # A widened [M, O] bias is tiled exactly like the output. A broadcast + # bias is not: it only has the trailing O extent, so asking for the + # output's cube would have the DMA read M * O elements out of an + # O-element buffer. That is invisible when everything already sits in + # L2 and there is no transfer, and it is why this only ever showed up + # on the L3 and promote configurations. + biasShape = ctxt.lookup(operatorRepresentation['C']).shape + if len([d for d in biasShape if d != 1]) <= 1: + leading = len(biasShape) - 1 + CCube = HyperRectangle(tuple([0] * leading + [OOffset]), tuple([1] * leading + [OSize])) + else: + CCube = HyperRectangle(tuple(cube.offset), tuple(cube.dims)) inputAddCubes.append(CCube) inputLoadSchedule = [] diff --git a/Deeploy/Targets/PULPOpen/Tiler.py b/Deeploy/Targets/PULPOpen/Tiler.py index d8d31a6b..004cbbfd 100644 --- a/Deeploy/Targets/PULPOpen/Tiler.py +++ b/Deeploy/Targets/PULPOpen/Tiler.py @@ -32,6 +32,7 @@ from Deeploy.Targets.PULPOpen.TileConstraints.AveragePoolTileConstraint import AveragePoolCTileConstraint from Deeploy.Targets.PULPOpen.TileConstraints.BatchNormTileConstraint import BatchNormalizationGradTileConstraint, \ BatchNormInternalTileConstraint +from Deeploy.Targets.PULPOpen.TileConstraints.BroadcastAddTileConstraint import PULPBroadcastAddTileConstraint from Deeploy.Targets.PULPOpen.TileConstraints.ConvGradConstraint import ConvGradBTileConstraint, \ ConvGradW2DTileConstraint, ConvGradX2DIm2ColHWTileConstraint, DWConvGradW2DTileConstraint, \ DWConvGradX2DTileConstraint, PWConvGradWTileConstraint, PWConvGradXTileConstraint @@ -159,7 +160,7 @@ tileConstraint = UnaryTileConstraint()) PULPAddTilingReadyBindings = TilingReadyNodeBindings(nodeBindings = PULPAddBindings, - tileConstraint = AddTileConstraint()) + tileConstraint = PULPBroadcastAddTileConstraint()) PULPSoftmaxTilingReadyBindings = TilingReadyNodeBindings(nodeBindings = PULPSoftmaxBindings, tileConstraint = iSoftmaxTileConstraint()) diff --git a/Deeploy/Targets/PULPOpen/TopologyOptimizationPasses/Passes.py b/Deeploy/Targets/PULPOpen/TopologyOptimizationPasses/Passes.py index 67179e12..b2296dc3 100644 --- a/Deeploy/Targets/PULPOpen/TopologyOptimizationPasses/Passes.py +++ b/Deeploy/Targets/PULPOpen/TopologyOptimizationPasses/Passes.py @@ -344,3 +344,126 @@ def __init__(self): name = "_MERGE_GEMM_MATMUL_RQ_PASS" super().__init__(graph, _merge_gemm_rq_fun, name) + + +@contextagnostic +class FoldDequantIntoMatMulPass(Pass): + """Fold a weight's Dequant into the MatMul or Gemm nodes that consume it. + + Weight-only quantisation stores a frozen weight as int8 and needs it as float + for an fp32 matmul. Left as its own node, the Dequant materialises the whole + dequantised matrix, which then has to live from the forward pass to the + backward one or be recomputed -- on CCT-QLoRA, 1088 KB of fp32 tensors whose + Dequant nodes account for 239.95M cycles, only 48.47M of which is the kernel. + QLoRA instead dequantises inside the matmul and discards the value; + PULP_MatMul_fp32_i8_fp32_unroll1x7 does that, given the per-tensor scale and + zero point this pass attaches. + + Written as a plain graph walk rather than a pattern replacement. The Dequant + worth folding is read by TWO nodes, the forward matmul and its gradient, and + that fan-out is what the sequential matcher skips; a Dequant with one consumer + is the cheap tail. Each consumer then dequantises independently, which repeats + the conversion on purpose -- it is register-level work inside the kernel, and + cheaper than keeping a materialised fp32 copy alive across the step. + + Only a win alongside constant deduplication: removing the Dequant gives the + int8 constant two direct consumers, and _duplicateConstants stores one copy per + consumer unless byte-identical constants are shared. + """ + + # Conv is here because at this stage -- before parsing inserts the NCHW->NHWC + # transposes -- the tokenizer's quantised weight reaches its Conv directly. + # Reading the post-parsing graph suggests otherwise and sent three attempts at + # this after a Transpose that does not exist yet when the pass runs. + FOLDABLE_OPS = ('MatMul', 'Gemm', 'Conv') + + def run_pass(self, graph: gs.Graph): + folded = 0 + for node in list(graph.nodes): + if node.op != 'Dequant' or not node.outputs or node.outputs[0] is None: + continue + dequantOut = node.outputs[0] + consumers = list(dequantOut.outputs) + if not consumers: + continue + # Every consumer must take it as the weight operand, which is the one + # the kernels dequantise. A Dequant feeding anything else is left alone. + if not all( + c.op in self.FOLDABLE_OPS and len(c.inputs) >= 2 and c.inputs[1] is dequantOut for c in consumers): + continue + quantised = node.inputs[0] if node.inputs else None + if quantised is None: + continue + + # The stored array is int8 but the tensor can carry a wider dtype from + # the exporter; with the Dequant in place that never mattered, because + # its output was explicitly float32. Once folded, this dtype is what the + # binding checker sees, and a wrong one silently selects the fp32 kernel + # and reads the int8 weight as float -- which runs, reports no error, and + # computes nonsense four times faster. + # gs.Constant derives dtype from values and does not let it be set, so + # rewriting the array is both necessary and sufficient. A Variable that + # is not a Constant cannot be corrected here and is left for the fp32 + # path rather than folded into a kernel that would misread it. + if not isinstance(quantised, gs.Constant): + continue + quantised.values = np.asarray(quantised.values).astype(np.int8) + + # Each consumer takes the int8 tensor directly and carries the scale and + # zero point as attributes, so it dequantises inside the kernel. Several + # consumers each repeat that conversion on purpose: it is register-level + # work, and cheaper than keeping a materialised fp32 copy alive between a + # forward use and its backward one. + # Every consumer reads the same int8 constant. A ConstantBuffer is global + # and read-only, so sharing one across the forward MatMul and its backward + # Gemm is sound, and it is what keeps the memory win: giving each consumer + # its own copy charges the full constant again per clone, which on the int8 + # QLoRA CCT took persistent from 662 KB to 938 KB. + for consumer in consumers: + consumer.inputs[1] = quantised + consumer.attrs['dequant_scale'] = float(node.attrs.get('scale', 1.0)) + consumer.attrs['dequant_zero_point'] = int(node.attrs.get('zero_point', 0)) + + node.inputs.clear() + node.outputs.clear() + folded += 1 + + if folded: + graph.cleanup() + return graph + + +def _fold_dequant_fun(graph: gs.Graph, match: Match, name: str): + """Rewire a Dequant's consumer to read the int8 tensor and dequantise in-kernel.""" + matched = [m for k, m in match.nodes_map.items()] + dequant, consumer = matched[0], matched[1] + + dequantOut = dequant.outputs[0] + # Only the weight operand is folded: the kernels dequantise their second input. + if len(consumer.inputs) < 2 or consumer.inputs[1] is not dequantOut: + return graph + quantised = dequant.inputs[0] if dequant.inputs else None + if quantised is None: + return graph + + # Every consumer of this Dequant takes the int8 tensor and dequantises it + # itself. That repeats the conversion on purpose: it is register-level work + # inside the kernel, and cheaper than keeping a materialised fp32 copy alive + # between a forward use and the backward one that reads it again. + consumers = list(dequantOut.outputs) + if not all(len(c.inputs) >= 2 and c.inputs[1] is dequantOut for c in consumers): + return graph + + for c in consumers: + c.inputs[1] = quantised + c.attrs['dequant_scale'] = float(dequant.attrs.get('scale', 1.0)) + c.attrs['dequant_zero_point'] = int(dequant.attrs.get('zero_point', 0)) + + # deleteNode() reconnects a node's output to its input, which assumes the node + # is a pass-through. This one is not any more: its consumers now read the int8 + # tensor directly, so its output is an orphan. Detach it and let cleanup() + # collect the node instead. + dequant.inputs.clear() + dequant.outputs.clear() + graph.cleanup() + return graph diff --git a/Deeploy/Targets/PULPOpen/TypeCheckers.py b/Deeploy/Targets/PULPOpen/TypeCheckers.py index e3096241..1841633a 100644 --- a/Deeploy/Targets/PULPOpen/TypeCheckers.py +++ b/Deeploy/Targets/PULPOpen/TypeCheckers.py @@ -7,6 +7,7 @@ from Deeploy.AbstractDataTypes import Pointer from Deeploy.CommonExtensions.TypeCheckers.SignPropTypeChecker import SignPropTypeChecker from Deeploy.DeeployTypes import OperatorRepresentation, VariableBuffer +from Deeploy.Targets.Generic.TypeCheckers import ConvChecker, GEMMChecker, MatMulChecker class PULPDMASliceChecker(SignPropTypeChecker): @@ -136,3 +137,38 @@ def _inferSignedness(self, inputs: List[VariableBuffer], # Override this. This should compute the signednes of each output node of the Layer def checkOutputType(self, inputs: List[VariableBuffer], operatorRepresentation: OperatorRepresentation) -> bool: return True + + +class _DequantFoldedMixin: + """Match only a node the Dequant fold has rewritten. + + The int8 kernel variants must not be selected on type alone. A binding is chosen + before a weight constant's type is fixed, and typeInferGlobalCtxt then types that + constant from the SELECTED binding's input_types, so whichever of the fp32 and + int8 bindings is tried first claims the weight and defines its type. Ordering them + is therefore not a fix: int8-first retypes ordinary fp32 weights to int8 and + silently quantises them (FP32 Conv kernel test, 512 errors of 512), while + fp32-first retypes folded int8 weights upward and inflates them (CCT-QLoRA + persistent 409 KB -> 1094 KB, with the int8 weights landing as float64). + + FoldDequantIntoMatMulPass marks every node it rewrites with dequant_scale, so that + attribute, not the operand type, is what distinguishes the two cases. + """ + + def typeCheck(self, ctxt, node, operatorRepresentation) -> bool: + # The fold marks every node it rewrites; nothing else may select this binding. + if 'dequant_scale' not in node.attrs: + return ctxt, False + return super().typeCheck(ctxt, node, operatorRepresentation) + + +class PULPDequantMatMulChecker(_DequantFoldedMixin, MatMulChecker): + pass + + +class PULPDequantGEMMChecker(_DequantFoldedMixin, GEMMChecker): + pass + + +class PULPDequantConvChecker(_DequantFoldedMixin, ConvChecker): + pass diff --git a/DeeployTest/Tests/Models/Training/CCT/cct_train/recompute_checkmate.json b/DeeployTest/Tests/Models/Training/CCT/cct_train/recompute_checkmate.json new file mode 100644 index 00000000..b59d8b5b --- /dev/null +++ b/DeeployTest/Tests/Models/Training/CCT/cct_train/recompute_checkmate.json @@ -0,0 +1 @@ +{"seq": [["node_1_tokenizer_conv_layers_conv_layers_0_conv_layers_0_0_Conv_Conv_input_0_transpose", 0], ["node_1_tokenizer_conv_layers_conv_layers_0_conv_layers_0_0_Conv_Conv", 0], ["node_1_tokenizer_conv_layers_conv_layers_0_conv_layers_0_0_Conv_Conv_node_0_tokenizer_conv_layers_conv_layers_0_conv_layers_0_0_Conv__0_tensor_pre_transpose", 0], ["node_2_tokenizer_conv_layers_conv_layers_0_conv_layers_0_1_Relu_Relu", 0], ["node_3_tokenizer_conv_layers_conv_layers_0_conv_layers_0_2_MaxPool_MaxPool_node_0_tokenizer_conv_layers_conv_layers_0_conv_layers_0_1_Relu__0_tensor_transpose", 0], ["node_3_tokenizer_conv_layers_conv_layers_0_conv_layers_0_2_MaxPool_MaxPool", 0], ["node_4_tokenizer_conv_layers_conv_layers_1_conv_layers_1_0_Conv_Conv", 0], ["node_4_tokenizer_conv_layers_conv_layers_1_conv_layers_1_0_Conv_Conv_node_0_tokenizer_conv_layers_conv_layers_1_conv_layers_1_0_Conv__0_tensor_pre_transpose", 0], ["node_5_tokenizer_conv_layers_conv_layers_1_conv_layers_1_1_Relu_Relu", 0], ["node_6_tokenizer_conv_layers_conv_layers_1_conv_layers_1_2_MaxPool_MaxPool_node_0_tokenizer_conv_layers_conv_layers_1_conv_layers_1_1_Relu__0_tensor_transpose", 0], ["node_6_tokenizer_conv_layers_conv_layers_1_conv_layers_1_2_MaxPool_MaxPool", 0], ["node_6_tokenizer_conv_layers_conv_layers_1_conv_layers_1_2_MaxPool_MaxPool_node_0_tokenizer_conv_layers_conv_layers_1_conv_layers_1_2_MaxPool__0_tensor_pre_transpose", 0], ["node_15_tokenizer_flattener_Reshape_Reshape", 0], ["node_16_tokenizer_Transpose_Transpose", 0], ["node_17_classifier_Add_Add", 0], ["node_21_classifier_blocks_0_pre_norm_LayerNormalization_LayerNormalization", 0], ["node_61_classifier_blocks_0_self_attn_v_proj_MatMul_MatMul", 0], ["node_62_classifier_blocks_0_self_attn_Reshape_2_Reshape", 0], ["node_26_classifier_blocks_0_self_attn_q_proj_MatMul_MatMul", 0], ["node_55_classifier_blocks_0_self_attn_Reshape_Reshape", 0], ["node_58_classifier_blocks_0_self_attn_k_proj_MatMul_MatMul", 0], ["node_59_classifier_blocks_0_self_attn_Reshape_1_Reshape", 0], ["node_63_classifier_blocks_0_self_attn_Transpose_1_Transpose_node_72_classifier_blocks_0_self_attn_MatMul_1_MatMul_transpose_in", 0], ["node_63_classifier_blocks_0_self_attn_Transpose_1_Transpose_node_72_classifier_blocks_0_self_attn_MatMul_1_GradFusedMatMul_0_1_Gemm_backward_transpose_in", 0], ["node_56_classifier_blocks_0_self_attn_Transpose_Transpose_node_65_classifier_blocks_0_self_attn_MatMul_MatMul_transpose_in", 0], ["node_56_classifier_blocks_0_self_attn_Transpose_Transpose_node_65_classifier_blocks_0_self_attn_MatMul_GradFusedMatMul_1_1_Gemm_backward_transpose_in", 0], ["node_64_classifier_blocks_0_self_attn_Transpose_2_Transpose_node_65_classifier_blocks_0_self_attn_MatMul_MatMul_transpose_in", 0], ["node_65_classifier_blocks_0_self_attn_MatMul_MatMul", 0], ["node_64_classifier_blocks_0_self_attn_Transpose_2_Transpose_node_65_classifier_blocks_0_self_attn_MatMul_GradFusedMatMul_0_1_Gemm_backward_transpose_in", 0], ["node_67_classifier_blocks_0_self_attn_Mul_Mul", 0], ["node_68_classifier_blocks_0_self_attn_Softmax_Softmax", 0], ["node_72_classifier_blocks_0_self_attn_MatMul_1_MatMul", 0], ["node_73_classifier_blocks_0_self_attn_Transpose_3_Transpose", 0], ["node_81_classifier_blocks_0_self_attn_Reshape_3_Reshape", 0], ["node_84_classifier_blocks_0_self_attn_proj_Add_Gemm", 0], ["node_88_classifier_blocks_0_Add_Add", 0], ["node_89_classifier_blocks_0_norm1_LayerNormalization_LayerNormalization", 0], ["node_92_classifier_blocks_0_linear1_Add_Gemm", 0], ["BiasGelu_Add_Add", 0], ["BiasGelu_Gelu_Gelu", 0], ["node_107_classifier_blocks_0_linear2_Add_Gemm", 0], ["node_111_classifier_blocks_0_Add_3_Add", 0], ["node_112_classifier_blocks_1_pre_norm_LayerNormalization_LayerNormalization", 0], ["node_152_classifier_blocks_1_self_attn_v_proj_MatMul_MatMul", 0], ["node_153_classifier_blocks_1_self_attn_Reshape_2_Reshape", 0], ["node_117_classifier_blocks_1_self_attn_q_proj_MatMul_MatMul", 0], ["node_146_classifier_blocks_1_self_attn_Reshape_Reshape", 0], ["node_149_classifier_blocks_1_self_attn_k_proj_MatMul_MatMul", 0], ["node_150_classifier_blocks_1_self_attn_Reshape_1_Reshape", 0], ["node_154_classifier_blocks_1_self_attn_Transpose_1_Transpose_node_163_classifier_blocks_1_self_attn_MatMul_1_MatMul_transpose_in", 0], ["node_154_classifier_blocks_1_self_attn_Transpose_1_Transpose_node_163_classifier_blocks_1_self_attn_MatMul_1_GradFusedMatMul_0_1_Gemm_backward_transpose_in", 0], ["node_147_classifier_blocks_1_self_attn_Transpose_Transpose_node_156_classifier_blocks_1_self_attn_MatMul_MatMul_transpose_in", 0], ["node_147_classifier_blocks_1_self_attn_Transpose_Transpose_node_156_classifier_blocks_1_self_attn_MatMul_GradFusedMatMul_1_1_Gemm_backward_transpose_in", 0], ["node_155_classifier_blocks_1_self_attn_Transpose_2_Transpose_node_156_classifier_blocks_1_self_attn_MatMul_MatMul_transpose_in", 0], ["node_156_classifier_blocks_1_self_attn_MatMul_MatMul", 0], ["node_155_classifier_blocks_1_self_attn_Transpose_2_Transpose_node_156_classifier_blocks_1_self_attn_MatMul_GradFusedMatMul_0_1_Gemm_backward_transpose_in", 0], ["node_158_classifier_blocks_1_self_attn_Mul_Mul", 0], ["node_159_classifier_blocks_1_self_attn_Softmax_Softmax", 0], ["node_163_classifier_blocks_1_self_attn_MatMul_1_MatMul", 0], ["node_164_classifier_blocks_1_self_attn_Transpose_3_Transpose", 0], ["node_172_classifier_blocks_1_self_attn_Reshape_3_Reshape", 0], ["node_175_classifier_blocks_1_self_attn_proj_Add_Gemm", 0], ["node_179_classifier_blocks_1_Add_Add", 0], ["node_180_classifier_blocks_1_norm1_LayerNormalization_LayerNormalization", 0], ["node_183_classifier_blocks_1_linear1_Add_Gemm", 0], ["BiasGelu_token_1_Add_Add", 0], ["BiasGelu_token_1_Gelu_Gelu", 0], ["node_198_classifier_blocks_1_linear2_Add_Gemm", 0], ["node_202_classifier_blocks_1_Add_3_Add", 0], ["node_197_classifier_blocks_1_linear2_MatMul_GradReshape_3_Reshape_backward", 0], ["node_203_classifier_norm_LayerNormalization_LayerNormalization", 0], ["Transpose_Reshape_Reshape", 0], ["node_206_classifier_attention_pool_Add_Gemm", 0], ["Reshape_before_softmax_node_0_classifier_Softmax__0_Reshape", 0], ["Softmax_optimized_node_0_classifier_Softmax__0_Softmax", 0], ["Reshape_after_softmax_node_0_classifier_Softmax__0_Reshape", 0], ["Transpose_Reshape_token_3_Reshape", 0], ["node_209_classifier_MatMul_MatMul", 0], ["node_211_classifier_Squeeze_Squeeze", 0], ["node_212_classifier_fc_Gemm_Gemm", 0], ["onnxSoftmaxCrossEntropyLoss4_SoftmaxCrossEntropyLoss", 0], ["onnxSoftmaxCrossEntropyLoss4_GradSoftmaxCrossEntropyLossGrad_0_SoftmaxCrossEntropyLossGrad_backward", 0], ["node_212_classifier_fc_Gemm_GradReduceSum_3_ReduceSum_backward", 0], ["GradientAccumulator35_InPlaceAccumulatorV2_backward", 0], ["node_212_classifier_fc_Gemm_GradGemm_0_Gemm_backward", 0], ["node_211_classifier_Squeeze_GradUnsqueeze_0_Unsqueeze_backward", 0], ["node_209_classifier_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["node_205_classifier_attention_pool_MatMul_GradReshape_3_Reshape_backward", 0], ["Transpose_Reshape_token_3_GradReshape_1_Reshape_backward", 0], ["node_207_classifier_Softmax_GradSoftmaxGrad_13_0_SoftmaxGrad_backward", 0], ["node_206_classifier_attention_pool_Add_GradReduceSum_1_ReduceSum_backward", 0], ["GradientAccumulator3_InPlaceAccumulatorV2_backward", 0], ["node_205_classifier_attention_pool_MatMul_GradReshape_4_Reshape_backward", 0], ["node_205_classifier_attention_pool_MatMul_GradGemm_5_Gemm_backward", 0], ["Transpose_Reshape_GradReshape_1_Reshape_backward", 0], ["GradientAccumulator2_InPlaceAccumulatorV2_backward", 0], ["node_212_classifier_fc_Gemm_GradGemm_1_Gemm_backward", 0], ["GradientAccumulator34_InPlaceAccumulatorV2_backward", 0], ["node_209_classifier_MatMul_GradFusedMatMul_1_1_Gemm_backward", 0], ["node_205_classifier_attention_pool_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["AccumulateGrad_node_0_classifier_norm_LayerNormalization__0_grad_add_Add_backward", 0], ["node_203_classifier_norm_LayerNormalization_LayerNormalization", 1], ["node_203_classifier_norm_LayerNormalization_GradLayerNormalizationGrad_0_LayerNormalizationGrad_backward", 0], ["GradientAccumulator32_InPlaceAccumulatorV2_backward", 0], ["GradientAccumulator33_InPlaceAccumulatorV2_backward", 0], ["node_198_classifier_blocks_1_linear2_Add_GradReduceSum_1_ReduceSum_backward", 0], ["GradientAccumulator31_InPlaceAccumulatorV2_backward", 0], ["node_197_classifier_blocks_1_linear2_MatMul_GradReshape_4_Reshape_backward", 0], ["node_197_classifier_blocks_1_linear2_MatMul_GradGemm_5_Gemm_backward", 0], ["node_196_classifier_blocks_1_linear2_Transpose_GradTranspose_0_Transpose_backward", 0], ["GradientAccumulator30_InPlaceAccumulatorV2_backward", 0], ["node_197_classifier_blocks_1_linear2_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["node_21_classifier_blocks_0_pre_norm_LayerNormalization_LayerNormalization", 1], ["GeluGrad_105_GeluGrad_backward", 0], ["node_89_classifier_blocks_0_norm1_LayerNormalization_LayerNormalization", 1], ["BiasGelu_token_1_GradReduceSum_2_ReduceSum_backward", 0], ["GradientAccumulator18_InPlaceAccumulatorV2_backward", 0], ["node_183_classifier_blocks_1_linear1_Add_GradReduceSum_1_ReduceSum_backward", 0], ["GradientAccumulator27_InPlaceAccumulatorV2_backward", 0], ["node_182_classifier_blocks_1_linear1_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["AccumulateGrad_node_0_classifier_blocks_1_norm1_LayerNormalization__0_grad_add_Add_backward", 0], ["node_180_classifier_blocks_1_norm1_LayerNormalization_LayerNormalization", 1], ["node_180_classifier_blocks_1_norm1_LayerNormalization_GradLayerNormalizationGrad_0_LayerNormalizationGrad_backward", 0], ["GradientAccumulator28_InPlaceAccumulatorV2_backward", 0], ["GradientAccumulator29_InPlaceAccumulatorV2_backward", 0], ["node_182_classifier_blocks_1_linear1_MatMul_GradReshape_4_Reshape_backward", 0], ["node_182_classifier_blocks_1_linear1_MatMul_GradReshape_3_Reshape_backward", 0], ["node_182_classifier_blocks_1_linear1_MatMul_GradGemm_5_Gemm_backward", 0], ["node_181_classifier_blocks_1_linear1_Transpose_GradTranspose_0_Transpose_backward", 0], ["GradientAccumulator26_InPlaceAccumulatorV2_backward", 0], ["node_175_classifier_blocks_1_self_attn_proj_Add_GradReduceSum_1_ReduceSum_backward", 0], ["GradientAccumulator25_InPlaceAccumulatorV2_backward", 0], ["node_174_classifier_blocks_1_self_attn_proj_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["node_172_classifier_blocks_1_self_attn_Reshape_3_GradReshape_1_Reshape_backward", 0], ["node_174_classifier_blocks_1_self_attn_proj_MatMul_GradReshape_4_Reshape_backward", 0], ["node_174_classifier_blocks_1_self_attn_proj_MatMul_GradReshape_3_Reshape_backward", 0], ["node_174_classifier_blocks_1_self_attn_proj_MatMul_GradGemm_5_Gemm_backward", 0], ["node_173_classifier_blocks_1_self_attn_proj_Transpose_GradTranspose_0_Transpose_backward", 0], ["GradientAccumulator24_InPlaceAccumulatorV2_backward", 0], ["node_164_classifier_blocks_1_self_attn_Transpose_3_GradTranspose_0_Transpose_backward_node_163_classifier_blocks_1_self_attn_MatMul_1_GradFusedMatMul_1_1_Gemm_backward_transpose_in", 0], ["node_164_classifier_blocks_1_self_attn_Transpose_3_GradTranspose_0_Transpose_backward_node_163_classifier_blocks_1_self_attn_MatMul_1_GradFusedMatMul_0_1_Gemm_backward_transpose_in", 0], ["node_163_classifier_blocks_1_self_attn_MatMul_1_GradFusedMatMul_0_1_Gemm_backward", 0], ["node_163_classifier_blocks_1_self_attn_MatMul_1_GradFusedMatMul_1_1_Gemm_backward", 0], ["node_159_classifier_blocks_1_self_attn_Softmax_GradSoftmaxGrad_13_0_SoftmaxGrad_backward", 0], ["node_154_classifier_blocks_1_self_attn_Transpose_1_GradTranspose_0_Transpose_backward", 0], ["node_153_classifier_blocks_1_self_attn_Reshape_2_GradReshape_1_Reshape_backward", 0], ["node_112_classifier_blocks_1_pre_norm_LayerNormalization_LayerNormalization", 1], ["node_158_classifier_blocks_1_self_attn_Mul_GradMul_0_Mul_backward", 0], ["node_156_classifier_blocks_1_self_attn_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["node_156_classifier_blocks_1_self_attn_MatMul_GradFusedMatMul_1_1_Gemm_backward", 0], ["node_147_classifier_blocks_1_self_attn_Transpose_GradTranspose_0_Transpose_backward", 0], ["node_155_classifier_blocks_1_self_attn_Transpose_2_GradTranspose_0_Transpose_backward", 0], ["node_146_classifier_blocks_1_self_attn_Reshape_GradReshape_1_Reshape_backward", 0], ["node_150_classifier_blocks_1_self_attn_Reshape_1_GradReshape_1_Reshape_backward", 0], ["node_152_classifier_blocks_1_self_attn_v_proj_MatMul_GradReshape_4_Reshape_backward", 0], ["node_152_classifier_blocks_1_self_attn_v_proj_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["node_152_classifier_blocks_1_self_attn_v_proj_MatMul_GradReshape_3_Reshape_backward", 0], ["node_152_classifier_blocks_1_self_attn_v_proj_MatMul_GradGemm_5_Gemm_backward", 0], ["node_151_classifier_blocks_1_self_attn_v_proj_Transpose_GradTranspose_0_Transpose_backward", 0], ["GradientAccumulator23_InPlaceAccumulatorV2_backward", 0], ["node_117_classifier_blocks_1_self_attn_q_proj_MatMul_GradReshape_4_Reshape_backward", 0], ["node_117_classifier_blocks_1_self_attn_q_proj_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["node_117_classifier_blocks_1_self_attn_q_proj_MatMul_GradReshape_3_Reshape_backward", 0], ["BiasGelu_Gelu_Gelu", 1], ["node_117_classifier_blocks_1_self_attn_q_proj_MatMul_GradGemm_5_Gemm_backward", 0], ["node_116_classifier_blocks_1_self_attn_q_proj_Transpose_GradTranspose_0_Transpose_backward", 0], ["GradientAccumulator21_InPlaceAccumulatorV2_backward", 0], ["node_149_classifier_blocks_1_self_attn_k_proj_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["AccumulateGrad_node_0_classifier_blocks_1_pre_norm_LayerNormalization__0_grad_add_0_Add_backward", 0], ["AccumulateGrad_node_0_classifier_blocks_1_pre_norm_LayerNormalization__0_grad_add_1_Add_backward", 0], ["node_112_classifier_blocks_1_pre_norm_LayerNormalization_GradLayerNormalizationGrad_0_LayerNormalizationGrad_backward", 0], ["AccumulateGrad_node_0_classifier_blocks_0_Add_3__0_grad_add_Add_backward", 0], ["GradientAccumulator19_InPlaceAccumulatorV2_backward", 0], ["GradientAccumulator20_InPlaceAccumulatorV2_backward", 0], ["node_149_classifier_blocks_1_self_attn_k_proj_MatMul_GradReshape_4_Reshape_backward", 0], ["node_149_classifier_blocks_1_self_attn_k_proj_MatMul_GradReshape_3_Reshape_backward", 0], ["node_149_classifier_blocks_1_self_attn_k_proj_MatMul_GradGemm_5_Gemm_backward", 0], ["node_148_classifier_blocks_1_self_attn_k_proj_Transpose_GradTranspose_0_Transpose_backward", 0], ["GradientAccumulator22_InPlaceAccumulatorV2_backward", 0], ["node_107_classifier_blocks_0_linear2_Add_GradReduceSum_1_ReduceSum_backward", 0], ["GradientAccumulator17_InPlaceAccumulatorV2_backward", 0], ["node_106_classifier_blocks_0_linear2_MatMul_GradReshape_4_Reshape_backward", 0], ["node_106_classifier_blocks_0_linear2_MatMul_GradReshape_3_Reshape_backward", 0], ["node_106_classifier_blocks_0_linear2_MatMul_GradGemm_5_Gemm_backward", 0], ["node_105_classifier_blocks_0_linear2_Transpose_GradTranspose_0_Transpose_backward", 0], ["GradientAccumulator16_InPlaceAccumulatorV2_backward", 0], ["node_106_classifier_blocks_0_linear2_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["GeluGrad_148_GeluGrad_backward", 0], ["BiasGelu_GradReduceSum_2_ReduceSum_backward", 0], ["GradientAccumulator4_InPlaceAccumulatorV2_backward", 0], ["node_92_classifier_blocks_0_linear1_Add_GradReduceSum_1_ReduceSum_backward", 0], ["GradientAccumulator13_InPlaceAccumulatorV2_backward", 0], ["node_91_classifier_blocks_0_linear1_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["AccumulateGrad_node_0_classifier_blocks_0_norm1_LayerNormalization__0_grad_add_Add_backward", 0], ["node_89_classifier_blocks_0_norm1_LayerNormalization_GradLayerNormalizationGrad_0_LayerNormalizationGrad_backward", 0], ["GradientAccumulator14_InPlaceAccumulatorV2_backward", 0], ["GradientAccumulator15_InPlaceAccumulatorV2_backward", 0], ["node_91_classifier_blocks_0_linear1_MatMul_GradReshape_4_Reshape_backward", 0], ["node_91_classifier_blocks_0_linear1_MatMul_GradReshape_3_Reshape_backward", 0], ["node_91_classifier_blocks_0_linear1_MatMul_GradGemm_5_Gemm_backward", 0], ["node_90_classifier_blocks_0_linear1_Transpose_GradTranspose_0_Transpose_backward", 0], ["GradientAccumulator12_InPlaceAccumulatorV2_backward", 0], ["node_84_classifier_blocks_0_self_attn_proj_Add_GradReduceSum_1_ReduceSum_backward", 0], ["GradientAccumulator11_InPlaceAccumulatorV2_backward", 0], ["node_83_classifier_blocks_0_self_attn_proj_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["node_81_classifier_blocks_0_self_attn_Reshape_3_GradReshape_1_Reshape_backward", 0], ["node_83_classifier_blocks_0_self_attn_proj_MatMul_GradReshape_4_Reshape_backward", 0], ["node_83_classifier_blocks_0_self_attn_proj_MatMul_GradReshape_3_Reshape_backward", 0], ["node_83_classifier_blocks_0_self_attn_proj_MatMul_GradGemm_5_Gemm_backward", 0], ["node_82_classifier_blocks_0_self_attn_proj_Transpose_GradTranspose_0_Transpose_backward", 0], ["GradientAccumulator10_InPlaceAccumulatorV2_backward", 0], ["node_73_classifier_blocks_0_self_attn_Transpose_3_GradTranspose_0_Transpose_backward_node_72_classifier_blocks_0_self_attn_MatMul_1_GradFusedMatMul_1_1_Gemm_backward_transpose_in", 0], ["node_73_classifier_blocks_0_self_attn_Transpose_3_GradTranspose_0_Transpose_backward_node_72_classifier_blocks_0_self_attn_MatMul_1_GradFusedMatMul_0_1_Gemm_backward_transpose_in", 0], ["node_72_classifier_blocks_0_self_attn_MatMul_1_GradFusedMatMul_0_1_Gemm_backward", 0], ["node_72_classifier_blocks_0_self_attn_MatMul_1_GradFusedMatMul_1_1_Gemm_backward", 0], ["node_68_classifier_blocks_0_self_attn_Softmax_GradSoftmaxGrad_13_0_SoftmaxGrad_backward", 0], ["node_63_classifier_blocks_0_self_attn_Transpose_1_GradTranspose_0_Transpose_backward", 0], ["node_62_classifier_blocks_0_self_attn_Reshape_2_GradReshape_1_Reshape_backward", 0], ["node_67_classifier_blocks_0_self_attn_Mul_GradMul_0_Mul_backward", 0], ["node_65_classifier_blocks_0_self_attn_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["node_65_classifier_blocks_0_self_attn_MatMul_GradFusedMatMul_1_1_Gemm_backward", 0], ["node_56_classifier_blocks_0_self_attn_Transpose_GradTranspose_0_Transpose_backward", 0], ["node_64_classifier_blocks_0_self_attn_Transpose_2_GradTranspose_0_Transpose_backward", 0], ["node_55_classifier_blocks_0_self_attn_Reshape_GradReshape_1_Reshape_backward", 0], ["node_59_classifier_blocks_0_self_attn_Reshape_1_GradReshape_1_Reshape_backward", 0], ["node_61_classifier_blocks_0_self_attn_v_proj_MatMul_GradReshape_4_Reshape_backward", 0], ["node_61_classifier_blocks_0_self_attn_v_proj_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["node_61_classifier_blocks_0_self_attn_v_proj_MatMul_GradReshape_3_Reshape_backward", 0], ["node_61_classifier_blocks_0_self_attn_v_proj_MatMul_GradGemm_5_Gemm_backward", 0], ["node_60_classifier_blocks_0_self_attn_v_proj_Transpose_GradTranspose_0_Transpose_backward", 0], ["GradientAccumulator9_InPlaceAccumulatorV2_backward", 0], ["node_26_classifier_blocks_0_self_attn_q_proj_MatMul_GradReshape_4_Reshape_backward", 0], ["node_26_classifier_blocks_0_self_attn_q_proj_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["node_26_classifier_blocks_0_self_attn_q_proj_MatMul_GradReshape_3_Reshape_backward", 0], ["node_26_classifier_blocks_0_self_attn_q_proj_MatMul_GradGemm_5_Gemm_backward", 0], ["node_25_classifier_blocks_0_self_attn_q_proj_Transpose_GradTranspose_0_Transpose_backward", 0], ["GradientAccumulator7_InPlaceAccumulatorV2_backward", 0], ["node_58_classifier_blocks_0_self_attn_k_proj_MatMul_GradFusedMatMul_0_1_Gemm_backward", 0], ["AccumulateGrad_node_0_classifier_blocks_0_pre_norm_LayerNormalization__0_grad_add_0_Add_backward", 0], ["AccumulateGrad_node_0_classifier_blocks_0_pre_norm_LayerNormalization__0_grad_add_1_Add_backward", 0], ["node_21_classifier_blocks_0_pre_norm_LayerNormalization_GradLayerNormalizationGrad_0_LayerNormalizationGrad_backward", 0], ["AccumulateGrad_node_0_classifier_Add__0_grad_add_Add_backward", 0], ["GradientAccumulator1_InPlaceAccumulatorV2_backward", 0], ["GradientAccumulator5_InPlaceAccumulatorV2_backward", 0], ["GradientAccumulator6_InPlaceAccumulatorV2_backward", 0], ["node_58_classifier_blocks_0_self_attn_k_proj_MatMul_GradReshape_4_Reshape_backward", 0], ["node_58_classifier_blocks_0_self_attn_k_proj_MatMul_GradReshape_3_Reshape_backward", 0], ["node_58_classifier_blocks_0_self_attn_k_proj_MatMul_GradGemm_5_Gemm_backward", 0], ["node_57_classifier_blocks_0_self_attn_k_proj_Transpose_GradTranspose_0_Transpose_backward", 0], ["GradientAccumulator8_InPlaceAccumulatorV2_backward", 0]]} \ No newline at end of file diff --git a/DeeployTest/Tests/Models/Training/CCT_LoRA_R1/cct_lorar1_optimizer/network.onnx b/DeeployTest/Tests/Models/Training/CCT_LoRA_R1/cct_lorar1_optimizer/network.onnx new file mode 100644 index 00000000..1e34b95a Binary files /dev/null and b/DeeployTest/Tests/Models/Training/CCT_LoRA_R1/cct_lorar1_optimizer/network.onnx differ diff --git a/DeeployTest/Tests/Models/Training/CCT_LoRA_R1/cct_lorar1_train/inputs.npz b/DeeployTest/Tests/Models/Training/CCT_LoRA_R1/cct_lorar1_train/inputs.npz new file mode 100644 index 00000000..fb359eca Binary files /dev/null and b/DeeployTest/Tests/Models/Training/CCT_LoRA_R1/cct_lorar1_train/inputs.npz differ diff --git a/DeeployTest/Tests/Models/Training/CCT_LoRA_R1/cct_lorar1_train/network.onnx b/DeeployTest/Tests/Models/Training/CCT_LoRA_R1/cct_lorar1_train/network.onnx new file mode 100644 index 00000000..78d8cc3e Binary files /dev/null and b/DeeployTest/Tests/Models/Training/CCT_LoRA_R1/cct_lorar1_train/network.onnx differ diff --git a/DeeployTest/Tests/Models/Training/CCT_LoRA_R1/cct_lorar1_train/outputs.npz b/DeeployTest/Tests/Models/Training/CCT_LoRA_R1/cct_lorar1_train/outputs.npz new file mode 100644 index 00000000..addab2c9 Binary files /dev/null and b/DeeployTest/Tests/Models/Training/CCT_LoRA_R1/cct_lorar1_train/outputs.npz differ diff --git a/DeeployTest/Tests/Models/Training/CCT_QLORA_FT/cct_qlorar1_optimizer/network.onnx b/DeeployTest/Tests/Models/Training/CCT_QLORA_FT/cct_qlorar1_optimizer/network.onnx new file mode 100644 index 00000000..1e34b95a Binary files /dev/null and b/DeeployTest/Tests/Models/Training/CCT_QLORA_FT/cct_qlorar1_optimizer/network.onnx differ diff --git a/DeeployTest/Tests/Models/Training/CCT_QLORA_FT/cct_qlorar1_train/inputs.npz b/DeeployTest/Tests/Models/Training/CCT_QLORA_FT/cct_qlorar1_train/inputs.npz new file mode 100644 index 00000000..fb359eca Binary files /dev/null and b/DeeployTest/Tests/Models/Training/CCT_QLORA_FT/cct_qlorar1_train/inputs.npz differ diff --git a/DeeployTest/Tests/Models/Training/CCT_QLORA_FT/cct_qlorar1_train/network.onnx b/DeeployTest/Tests/Models/Training/CCT_QLORA_FT/cct_qlorar1_train/network.onnx new file mode 100644 index 00000000..66a4b42c Binary files /dev/null and b/DeeployTest/Tests/Models/Training/CCT_QLORA_FT/cct_qlorar1_train/network.onnx differ diff --git a/DeeployTest/Tests/Models/Training/CCT_QLORA_FT/cct_qlorar1_train/outputs.npz b/DeeployTest/Tests/Models/Training/CCT_QLORA_FT/cct_qlorar1_train/outputs.npz new file mode 100644 index 00000000..addab2c9 Binary files /dev/null and b/DeeployTest/Tests/Models/Training/CCT_QLORA_FT/cct_qlorar1_train/outputs.npz differ diff --git a/DeeployTest/conftest.py b/DeeployTest/conftest.py index 17a0c67f..5006a52b 100644 --- a/DeeployTest/conftest.py +++ b/DeeployTest/conftest.py @@ -89,6 +89,9 @@ def pytest_configure(config: pytest.Config) -> None: config.addinivalue_line("markers", "doublebuffer: mark test as double-buffer configuration") config.addinivalue_line("markers", "l2: mark test as L2 default memory level") config.addinivalue_line("markers", "l3: mark test as L3 default memory level") + config.addinivalue_line( + "markers", "recompute: mark test as replaying a gradient-checkpointing schedule " + "(activations regenerated before the backward pass instead of held live)") config.addinivalue_line("markers", "wmem: mark test as using Neureka weight memory") config.addinivalue_line("markers", "promote: mark test as exercising the L3->L2 tensor promotion pass") config.addinivalue_line("markers", "dma: mark test as DMA test") diff --git a/DeeployTest/deeployTrainingRunner_tiled_gap9.py b/DeeployTest/deeployTrainingRunner_tiled_gap9.py index beec9f66..9302066d 100644 --- a/DeeployTest/deeployTrainingRunner_tiled_gap9.py +++ b/DeeployTest/deeployTrainingRunner_tiled_gap9.py @@ -17,6 +17,12 @@ def setup_parser(parser): help = 'Keep forward convs channels-first (NCHW), binding the *_CHW kernels ' '(removes the NCHW<->NHWC transpose around every conv; required for ' 'MobileNetV1 training to fit GAP9 L1).\n') + parser.add_argument('--recomputeSchedule', + type = str, + default = None, + help = 'Path to a gradient-checkpointing schedule to replay on the deployed ' + 'graph, regenerating checkpointed activations before the backward pass reads ' + 'them instead of keeping them live.\n') parser.add_argument('--powerMeasurement', action = 'store_true', default = False, diff --git a/DeeployTest/testMVPTraining.py b/DeeployTest/testMVPTraining.py index 0c644af0..44630edd 100644 --- a/DeeployTest/testMVPTraining.py +++ b/DeeployTest/testMVPTraining.py @@ -15,7 +15,8 @@ from testUtils.testRunner import TestGeneratorArgumentParser from testUtils.tilingUtils import TrainingDBOnlyL3Tiler, TrainingDBTiler, TrainingSBTiler from testUtils.trainingUtils import _GRAD_ACC, _infer_data_size, _infer_n_accum, _infer_num_data_inputs, \ - _infer_total_mb, _load_reference_losses, _memoryMinimisingScheduler, add_training_inference_args + _infer_total_mb, _load_reference_losses, _memoryMinimisingScheduler, add_training_inference_args, \ + recomputeScheduler from testUtils.typeMapping import inferTypeAndOffset from Deeploy.AbstractDataTypes import PointerClass @@ -115,7 +116,8 @@ def generateTiledTrainingNetwork(args) -> None: deeployStateDir = _DEEPLOYSTATEDIR, inputOffsets = inputOffsets, conv_channels_first = args.convChannelsFirst, - scheduler = _memoryMinimisingScheduler) + scheduler = recomputeScheduler(args.recomputeSchedule) + if args.recomputeSchedule else _memoryMinimisingScheduler) # 7. Set up memory hierarchy. L3 = MemoryLevel(name = "L3", neighbourNames = ["L2"], size = 64_000_000) @@ -314,6 +316,13 @@ def generateTiledTrainingNetwork(args) -> None: type = int, default = 131072, help = 'Bytes reserved in L2 for tile staging') + parser.add_argument('--recomputeSchedule', + type = str, + default = None, + help = 'Path to a recompute schedule: JSON {"seq": [[node, is_recompute], ...]} in run ' + 'order, as emitted by a gradient-checkpointing solve. Replays it on the deployed graph, ' + 'regenerating each checkpointed activation just before the backward pass reads it ' + 'instead of keeping it live. Without this the memory-minimising schedule is used.') parser.add_argument('--convChannelsFirst', action = 'store_true', help = 'Use channels-first forward convolutions (no NCHW<->NHWC transpose).') diff --git a/DeeployTest/testUtils/deeployTrainingRunner.py b/DeeployTest/testUtils/deeployTrainingRunner.py index d15c22af..627057dd 100644 --- a/DeeployTest/testUtils/deeployTrainingRunner.py +++ b/DeeployTest/testUtils/deeployTrainingRunner.py @@ -122,6 +122,8 @@ def main(tiling_enabled: bool = False, gen_args.extend(['--input-offset-map'] + list(args.input_offset_map)) if getattr(args, 'convChannelsFirst', False): gen_args.append('--convChannelsFirst') + if getattr(args, 'recomputeSchedule', None): + gen_args.append(f'--recomputeSchedule={args.recomputeSchedule}') if tiling_enabled: if getattr(args, 'defaultMemLevel', None): diff --git a/DeeployTest/testUtils/pytestRunner.py b/DeeployTest/testUtils/pytestRunner.py index 5750d51b..ddfb8b09 100644 --- a/DeeployTest/testUtils/pytestRunner.py +++ b/DeeployTest/testUtils/pytestRunner.py @@ -64,6 +64,7 @@ def create_test_config( training_num_data_inputs: Optional[int] = None, training_tolerance: Optional[float] = None, training_conv_channels_first: bool = False, + training_recompute_schedule: str = None, promote_to_l2: bool = False, promote_to_l2_strategy: str = "cycle-aware", promote_to_l2_headroom: int = 131072, @@ -128,6 +129,9 @@ def create_test_config( if training and training_conv_channels_first: gen_args_list.append("--convChannelsFirst") + if training and training_recompute_schedule is not None: + gen_args_list.append(f"--recomputeSchedule={training_recompute_schedule}") + config = DeeployTestConfig( test_name = test_name_clean, test_dir = test_dir_abs, diff --git a/DeeployTest/testUtils/trainingUtils.py b/DeeployTest/testUtils/trainingUtils.py index 39cd940c..4b78bd63 100644 --- a/DeeployTest/testUtils/trainingUtils.py +++ b/DeeployTest/testUtils/trainingUtils.py @@ -30,7 +30,7 @@ import subprocess import sys from pathlib import Path -from typing import List, Optional +from typing import Dict, List, Optional, Set, Tuple import numpy as np import onnx_graphsurgeon as gs @@ -138,6 +138,209 @@ def _mockScheduler(graph: gs.Graph) -> List[List[gs.Node]]: return [[node] for node in graph.nodes] +# Suffix marking a node, tensor or constant created as a recompute clone. Also the +# signal that a graph has already been injected, so the second scheduler call +# replays the cached order instead of cloning a second time. +_CLONE_MARKER = "__ck" + +# Fraction of a schedule's entries that must name a node in the deployed graph before +# the replay is trusted. Not 100%: lowering legitimately renames a few nodes. +_MIN_SCHEDULE_COVERAGE = 0.9 + + +def _loadRecomputeSchedule(path: str) -> List[Tuple[str, bool]]: + """Read a recompute schedule: a list of (node name, is_recompute) in run order. + + The file is what a Checkmate solve emits, one entry per execution rather than + per node, so a tensor that is recomputed appears more than once. The first + appearance of a name is its original execution; every later one asks for a + fresh clone. + """ + with open(path) as handle: + data = json.load(handle) + return [(str(name), bool(isRecompute)) for name, isRecompute in data["seq"]] + + +def _liveOutputs(node: gs.Node) -> int: + return len([out for out in node.outputs if out is not None]) + + +def _nodesByName(graph: gs.Graph) -> Dict[str, gs.Node]: + """Map name to node, preferring one that still produces something. + + Lowering can leave a stale same-named node whose outputs have been cleared. A + plain comprehension may pick that one, and scheduling a zero-output node puts + an empty layer into layerBinding, which fails with an IndexError far from here. + """ + byName: Dict[str, gs.Node] = {} + for node in graph.nodes: + if node.name not in byName or (_liveOutputs(node) and not _liveOutputs(byName[node.name])): + byName[node.name] = node + return byName + + +def _spliceMissingNodes(schedule: List[gs.Node], graph: gs.Graph) -> List[gs.Node]: + """Insert nodes the schedule does not mention after their last producer. + + The scheduler runs twice, once for binding and once for the tiler, and lowering + between the two can add or rename nodes. Dropping them would leave their outputs + unproduced. Splicing each one in after its last scheduled predecessor keeps the + rest of the order byte-identical, which a full re-topologisation would not. + """ + placed = {id(node) for node in schedule} + for node in [n for n in graph.nodes if id(n) not in placed]: + producedNames = {t.name for t in node.inputs if t is not None and t.name} + insertAt = 0 + for index, scheduled in enumerate(schedule): + if any(out is not None and out.name in producedNames for out in scheduled.outputs): + insertAt = index + 1 + schedule.insert(insertAt, node) + return schedule + + +def _pruneDeadRecomputes(schedule: List[gs.Node], graph: gs.Graph) -> Tuple[List[gs.Node], int]: + """Drop scheduled nodes nothing reads, to a fixpoint. + + A solve may place a recompute in a stage whose consumers linearise before the + clone, so nothing is rewired to it. It may equally rewire every consumer of an + original to a clone, which makes the "recompute" a move and leaves the original + unread. Either way cleanup strips the output and the empty layer reaches + layerBinding. Dropping such a node is strictly better: less compute and less + memory. Removing one can orphan its producer, hence the fixpoint. + """ + graphOutputs = {out.name for out in graph.outputs if out is not None} + dropped = 0 + while True: + position = {id(node): index for index, node in enumerate(schedule)} + dead = [] + for node in schedule: + if any(out is not None and out.name in graphOutputs for out in node.outputs): + continue + isRead = any( + id(consumer) in position and position[id(consumer)] > position[id(node)] for out in node.outputs + if out is not None for consumer in out.outputs) + if not isRead: + dead.append(node) + if not dead: + return schedule, dropped + deadIds = {id(node) for node in dead} + schedule = [node for node in schedule if id(node) not in deadIds] + dropped += len(dead) + + +def recomputeScheduler(schedulePath: str, shareConstants: bool = True): + """Build a scheduler that replays a recompute schedule on the deployed graph. + + Gradient checkpointing trades compute for memory: an activation is dropped after + its forward use and regenerated before the backward one reads it. The decision of + what to drop comes from a solve over the training graph; this replays that + decision on the graph Deeploy actually deploys, which is not the graph the solve + saw, because lowering has run in between. + + Each recompute becomes a short-lived clone, and consumers scheduled after it read + the clone rather than the original, so the original's output dies at its last + forward use instead of staying live across the backward pass. + + ``shareConstants`` lets a clone reference the same constant as the original + rather than carrying a copy. A copy charges the full constant again for every + clone, so recomputing a node that reads a weight grows the persistent region + instead of shrinking the peak: measured on the int8 QLoRA CCT, 117 clones pushed + persistent from 662 KB to 938 KB. Sharing is sound because a ConstantBuffer is + global and read-only. + + The returned callable is used twice on the same graph, once for binding and once + for the tiler. The second call replays the order cached from the first: cloning + again would produce a second set of clones for the same recomputes. + """ + sequence = _loadRecomputeSchedule(schedulePath) + cachedOrder: List[Optional[List[str]]] = [None] + + def schedule(graph: gs.Graph) -> List[List[gs.Node]]: + byName = _nodesByName(graph) + + if cachedOrder[0] is not None and any(_CLONE_MARKER in node.name for node in graph.nodes): + replayed = [byName[name] for name in cachedOrder[0] if name in byName] + replayed = _spliceMissingNodes(replayed, graph) + return [[node] for node in replayed if _liveOutputs(node)] + + # A schedule is keyed on node names, and lowering renames nodes. If a solve is + # replayed against a graph it was not produced from, most names miss, every + # recompute silently vanishes, and what runs is the default order under a name + # that claims to be checkpointed. That degradation is invisible in a pass/fail + # CI result, so refuse it rather than report a green run of the wrong thing. + matched = sum(1 for name, _ in sequence if name in byName) + if matched < len(sequence) * _MIN_SCHEDULE_COVERAGE: + raise ValueError(f"Recompute schedule matches only {matched} of {len(sequence)} entries against this " + f"graph ({matched / max(len(sequence), 1):.0%}, need " + f"{_MIN_SCHEDULE_COVERAGE:.0%}). It was almost certainly produced from a different " + f"graph: replaying it would drop the recomputes and silently run the default order.") + + current: Dict[str, gs.Tensor] = {} + ordered: List[gs.Node] = [] + clones: List[gs.Node] = [] + alreadyRun: Set[str] = set() + instance = 0 + + for name, isRecompute in sequence: + original = byName.get(name) + if original is None: + continue + rewiredInputs = [current[t.name] if (t is not None and t.name in current) else t for t in original.inputs] + + if not isRecompute and name not in alreadyRun: + original.inputs = rewiredInputs + ordered.append(original) + for out in original.outputs: + if out.name: + current[out.name] = out + alreadyRun.add(name) + continue + + instance += 1 + cloneInputs = [] + for tensor in rewiredInputs: + if isinstance(tensor, gs.Constant) and not shareConstants: + cloneInputs.append( + gs.Constant(name = f"{tensor.name}{_CLONE_MARKER}{instance}", values = tensor.values)) + else: + cloneInputs.append(tensor) + cloneOutputs = [ + gs.Variable(name = f"{out.name}{_CLONE_MARKER}{instance}", dtype = out.dtype, shape = out.shape) + if out.name else out for out in original.outputs + ] + clone = gs.Node(op = original.op, + name = f"{original.name}{_CLONE_MARKER}{instance}", + attrs = dict(original.attrs), + inputs = cloneInputs, + outputs = cloneOutputs) + ordered.append(clone) + clones.append(clone) + for out, cloneOut in zip(original.outputs, cloneOutputs): + if out.name: + current[out.name] = cloneOut + + ordered, dropped = _pruneDeadRecomputes(ordered, graph) + if dropped: + log.info(f"[Recompute] dropped {dropped} recomputes nothing reads") + + liveClones = {id(node) for node in ordered} + graph.nodes.extend([node for node in clones if id(node) in liveClones]) + + # Nodes the schedule never named still have to run, and they cannot simply go + # at the end: a constant duplicated per consumer during frontend folding + # produces a tensor its consumer reads, so appending it after that consumer + # leaves the input unresolved and context lookup fails. Splice each one in + # after whatever produces its inputs. + ordered = _spliceMissingNodes(ordered, graph) + + cachedOrder[0] = [node.name for node in ordered] + log.info(f"[Recompute] {len(sequence)} scheduled executions, " + f"{len([n for n in ordered if _CLONE_MARKER in n.name])} recompute clones") + return [[node] for node in ordered if _liveOutputs(node)] + + return schedule + + # Element size assumed by _tensorSize. Deliberately uniform: this is the weight # of a tie-break heuristic, not an allocation, and a training graph is fp32 # almost everywhere. Reading tensor.dtype instead is not the free improvement it diff --git a/DeeployTest/test_gap9_tiled_config.py b/DeeployTest/test_gap9_tiled_config.py index f633be72..1078d930 100644 --- a/DeeployTest/test_gap9_tiled_config.py +++ b/DeeployTest/test_gap9_tiled_config.py @@ -91,6 +91,23 @@ "Models/Training/SimpleMLP/simplemlp_train": [64000], "Models/Training/Autoencoder/autoencoder_train": [128000], "Models/Training/DSCNN/dscnn_train": [128000, 64000], + # ResNet8 on-chip. It only fits channels-first: the NHWC path materialises a + # transposed copy of every conv weight -- 75 of them, 2950 KB in total -- because + # the pass that eliminates W^T covers Linear layers and not Conv. CHW kernels need + # no transpose at all, which takes the peak from 1410 KB to 1278 KB and brings it + # inside L2 with no recompute, no promotion and no double buffering. + # Measured at one mini-batch: Errors: 0, 43,531,388 cycles, L2_shared 164928 B. + # 116000, not 122000: CI trains four, and their accumulator buffers leave the L1 + # allocator 118 KB. + "Models/Training/ResNet8/resnet8_train": [116000], + # CCT-QLoRA on-chip. The frozen backbone is int8 and its Dequant is folded into + # the MatMul/Gemm/Conv, so the dequantised weights are never materialised: + # weight_sram is 48 KB and the arena needs 923 KB, which fits GAP9's real 1.5 MB + # L2 but not the 1000 KB runner default -- hence the l2 override below. + # 116000, not the 122000 the L3 entry uses: CI trains 4 mini-batches, whose + # accumulator buffers leave the L1 allocator 118 KB, and a 122000 arena does not + # fit alongside them. + "Models/Training/CCT_QLORA_FT/cct_qlorar1_train": [116000], } # L3 models: ResNet8, MobileNetV1, CCT exceed 1 MB L2 — weights spill to L3. @@ -102,12 +119,39 @@ "Models/Training/ResNet8/resnet8_train": [122000], "Models/Training/MobileNetV1/mobilenetv1_train": [116000], "Models/Training/CCT/cct_train": [122000], - "Models/Training/CCT_LoRA/cct_lora_train": [40000], + # Rank-4 LoRA on CCT-2 at the official spec (mlp_ratio=1), adapters on attention + # and FFN. 2292 KB peak, 54.5M cycles -- faster than full fine-tuning because the + # frozen base weights need no weight gradients. + "Models/Training/CCT_LoRA_R1/cct_lorar1_train": [122000], + # The same model with its frozen backbone quantised to int8. Exercises the folded + # Dequant path: without in-kernel dequantisation this model does not fit at all + # (minimalloc fails), and with it the weights reach the kernels as int8. + "Models/Training/CCT_QLORA_FT/cct_qlorar1_train": [122000], "Models/Training/SleepConViT/sleepconvit_train": [122000], "Models/Training/TSDR/tsdr_train": [122000], "Models/Training/MCUNet/mcunet_train": [116000], } +# Gradient-checkpointed training. Each entry replays a solved recompute schedule +# instead of keeping every forward activation live across the backward pass: an +# activation is dropped after its forward use and regenerated just before the +# gradient node reads it, trading cycles for peak memory. +# +# Only schedules verified end to end on gvsoc belong here. A schedule is keyed on +# node names, so one produced from a different graph would replay as the default +# order under a name claiming to be checkpointed; the replay refuses below 90% +# name coverage rather than report that as a pass. +# +# model schedule Errors cycles/step +# CCT (exact ILP) recompute_checkmate.json 0 68.63M vs 64.24M +# baseline (+6.8%) +L3_RECOMPUTE_TRAINING_MODELS = { + "Models/Training/CCT/cct_train": { + "l1": 122000, + "schedule": "Tests/Models/Training/CCT/cct_train/recompute_checkmate.json", + }, +} + # L3 double-buffered training. Only the L3<->L2 hop is double-buffered # (TrainingDBOnlyL3Tiler) so the L2 staging budget doesn't double. # ResNet8's strided ConvGrad InPlaceAccumulatorV2 is forced to coeff=1 (SB) @@ -161,6 +205,8 @@ "slave_stack": 512, # 0.80M vs 1.21M cyc/step (-33.7%) }, "Models/Training/ResNet8/resnet8_train": { + "conv_channels_first": True, # CHW convs; see the on-chip entry above + "l2": 1400000, # arena 1278 KB plus the static section, inside GAP9's 1.5 MB "cc_stack": 4096, # conv-light backward -> small CC stack, frees L1 for arena # arena 122000 + cc 4096 + slave 512*8 = 130192 < 131072 -> L1 stacks fit. # L1 vs L2 stacks: 47.8M vs 62.9M cyc/step (-23.9%, SB). With the gather @@ -195,8 +241,25 @@ # Also helps single-buffer (95.8M -> 75.3M). promote+DB+L1 is CCT's best. "slave_stack": 512, }, - "Models/Training/CCT_LoRA/cct_lora_train": { + "Models/Training/CCT_LoRA_R1/cct_lorar1_train": { "num_data_inputs": 1, + "tolerance": 5e-3, + "cc_stack": 4096, + "slave_stack": 512, + }, + "Models/Training/CCT_QLORA_FT/cct_qlorar1_train": { + "l2": 1572864, # GAP9's real 1.5 MB; the arena needs 923 KB of it + # arena 122000 + cc 4096 + slave 512*8 = 130192 < 131072, so the cluster + # stacks stay in L1. Without these the SDK defaults overflow it and gvsoc + # exits before producing any output. + "cc_stack": 4096, + "slave_stack": 512, + "num_data_inputs": 1, + # int8 weights dequantised in-kernel; the tolerance covers per-tensor + # symmetric quantisation of the frozen backbone, not a looser kernel. + "tolerance": 5e-2, + "cc_stack": 4096, + "slave_stack": 512, }, "Models/Training/SleepConViT/sleepconvit_train": { "num_data_inputs": 1, diff --git a/DeeployTest/test_platforms.py b/DeeployTest/test_platforms.py index 0559010c..16d2b03d 100644 --- a/DeeployTest/test_platforms.py +++ b/DeeployTest/test_platforms.py @@ -23,6 +23,7 @@ from test_gap9_tiled_config import L3_DOUBLEBUFFER_TRAINING_MODELS as GAP9_L3_DOUBLEBUFFER_TRAINING_MODELS from test_gap9_tiled_config import \ L3_DOUBLEBUFFER_TRAINING_PROMOTE_MODELS as GAP9_L3_DOUBLEBUFFER_TRAINING_PROMOTE_MODELS +from test_gap9_tiled_config import L3_RECOMPUTE_TRAINING_MODELS as GAP9_L3_RECOMPUTE_TRAINING_MODELS from test_gap9_tiled_config import L3_SINGLEBUFFER_MODELS as GAP9_L3_SINGLEBUFFER_MODELS from test_gap9_tiled_config import L3_SINGLEBUFFER_TRAINING_MODELS as GAP9_L3_SINGLEBUFFER_TRAINING_MODELS from test_gap9_tiled_config import \ @@ -1404,7 +1405,9 @@ def test_gap9_tiled_training_l2_singlebuffer(test_params, deeploy_test_dir, tool tiling = True, cores = GAP9_TILED_DEFAULT_CORES, l1 = l1, - l2 = 1024000, + # GAP9's L2 is 1.5 MB; 1024000 is the runner default. A model that needs + # the real capacity says so rather than being left out of the on-chip job. + l2 = overrides.get("l2", 1024000), default_mem_level = "L2", double_buffer = False, training = True, @@ -1462,6 +1465,56 @@ def test_gap9_tiled_training_l3_singlebuffer(test_params, deeploy_test_dir, tool run_and_assert_test(test_name, config, skipgen, skipsim) +@pytest.mark.gap9_tiled +@pytest.mark.training +@pytest.mark.singlebuffer +@pytest.mark.l3 +@pytest.mark.recompute +@pytest.mark.parametrize( + "test_name,recompute_config", + sorted(GAP9_L3_RECOMPUTE_TRAINING_MODELS.items()), + ids = lambda value: value.replace("/", "-") if isinstance(value, str) else "", +) +def test_gap9_tiled_training_l3_recompute(test_name, recompute_config, deeploy_test_dir, toolchain, toolchain_dir, + cmake_args, skipgen, skipsim) -> None: + """Train with a solved gradient-checkpointing schedule replayed on the deployed graph. + + This is the same model and budget as the singlebuffer case; what differs is that + checkpointed activations are regenerated before the backward pass reads them + rather than held live. It costs cycles and buys peak memory, so it is a separate + entry rather than a replacement, and both run. + """ + overrides = GAP9_TRAINING_MODEL_OVERRIDES.get(test_name, {}) + gap9_cmake_args = cmake_args + [f"NUM_CORES={GAP9_TILED_DEFAULT_CORES}"] + cc_stack = overrides.get("cc_stack") + if cc_stack is not None: + gap9_cmake_args = gap9_cmake_args + [f"CC_STACK_SIZE={cc_stack}"] + slave_stack = overrides.get("slave_stack") + if slave_stack is not None: + gap9_cmake_args = gap9_cmake_args + [f"SLAVESTACKSIZE={slave_stack}"] + config = create_test_config( + test_name = test_name, + platform = "GAP9", + simulator = "gvsoc", + deeploy_test_dir = deeploy_test_dir, + toolchain = toolchain, + toolchain_dir = toolchain_dir, + cmake_args = gap9_cmake_args, + tiling = True, + cores = GAP9_TILED_DEFAULT_CORES, + l1 = recompute_config["l1"], + l2 = 1024000, + default_mem_level = "L3", + double_buffer = False, + training = True, + training_num_data_inputs = overrides.get("num_data_inputs"), + training_tolerance = overrides.get("tolerance"), + training_conv_channels_first = overrides.get("conv_channels_first", False), + training_recompute_schedule = recompute_config["schedule"], + ) + run_and_assert_test(test_name, config, skipgen, skipsim) + + @pytest.mark.gap9_tiled @pytest.mark.training @pytest.mark.doublebuffer diff --git a/DeeployTest/test_siracusa_tiled_config.py b/DeeployTest/test_siracusa_tiled_config.py index 68c26d78..6476e111 100644 --- a/DeeployTest/test_siracusa_tiled_config.py +++ b/DeeployTest/test_siracusa_tiled_config.py @@ -212,12 +212,6 @@ # model artifacts. TRAINING_MODEL_OVERRIDES = { "Models/Training/CCT/cct_train": {}, - "Models/Training/CCT_LoRA/cct_lora_train": { - # Reduced from 32→4 mini-batches (2 optimizer steps, n_accum=2). - # Steps 0-3 are all within 2.5e-5 of ORT — no tolerance override needed. - # The old 32-step test compounded LoRA backward drift to ~1.2e-2 at - # step 27; 4 steps is sufficient coverage at default 1e-3 tolerance. - }, # conv_channels_first: run the convs natively in NCHW. In NHWC the conv # nets emit hundreds of NCHW<->NHWC Transpose ops (ResNet8 forward alone has # ~647) that stream through L3 HyperRAM — cheap on GAP9's memory but ~4x the diff --git a/TargetLibraries/GAP9/inc/kernel/GAP9Kernels.h b/TargetLibraries/GAP9/inc/kernel/GAP9Kernels.h index c3fac31f..773ca461 100644 --- a/TargetLibraries/GAP9/inc/kernel/GAP9Kernels.h +++ b/TargetLibraries/GAP9/inc/kernel/GAP9Kernels.h @@ -145,12 +145,28 @@ void PULP_DW_Conv2d_Im2Col_fp32_fp32_fp32_CHW( uint32_t pad_left, uint32_t pad_right, float32_t *__restrict__ pContextBuffer); +// MatMul, declared here for the same reason as the gradient kernels above: the +// generated code includes DeeployGAP9Math.h and never PULPOpen's +// kernel/Matmul.h, so without this the call is emitted against an implicit +// (int) declaration and the float arguments go under the wrong ABI. Kept in +// sync with TargetLibraries/PULPOpen/inc/kernel/Matmul.h. +void PULP_MatMul_fp32_fp32_fp32_unroll1x7(const float32_t *__restrict__ pSrcA, + const float32_t *__restrict__ pSrcB, + float32_t *__restrict__ pDstY, + uint32_t M, uint32_t N, uint32_t O); + +void PULP_MatMul_fp32_i8_fp32_unroll1x7(const float32_t *__restrict__ pSrcA, + const int8_t *__restrict__ pSrcB, + float32_t *__restrict__ pDstY, + uint32_t M, uint32_t N, uint32_t O, + float32_t scale, int32_t zeroPoint); + void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA, const float32_t *__restrict__ pSrcB, const float32_t *__restrict__ pDstC, float32_t *__restrict__ pDstY, uint32_t M, uint32_t N, uint32_t O, uint32_t transA, - uint32_t transB); + uint32_t transB, uint32_t biasStride); void PULP_GlobalAveragePoolGrad_fp32(const float32_t *dY, float32_t *dX, uint32_t N, uint32_t C, uint32_t H, diff --git a/TargetLibraries/PULPOpen/inc/kernel/Conv.h b/TargetLibraries/PULPOpen/inc/kernel/Conv.h index 6ae0ba99..d2c5d7c6 100644 --- a/TargetLibraries/PULPOpen/inc/kernel/Conv.h +++ b/TargetLibraries/PULPOpen/inc/kernel/Conv.h @@ -138,4 +138,15 @@ void PULP_DWConvGradX2d_fp32_fp32_fp32_CHW_tiled( uint32_t padding_y_bottom, // pad_right (unused here) uint16_t offset_grad_in_h, uint16_t offset_grad_in_w, uint16_t offset_grad_out_h, uint16_t offset_grad_out_w); + +// HWC im2col convolution with an int8 weight, dequantised in-kernel; see +// Convolution_fp32.c for why the scale factors out of the accumulation. +void PULP_Conv2d_Im2Col_fp32_i8_fp32_HWC( + const float32_t *__restrict__ pSrcA, uint32_t H, uint32_t W, uint32_t C, + const int8_t *__restrict__ pSrcB, uint32_t F_total, uint32_t P, uint32_t Q, + uint32_t SP, uint32_t SQ, const float32_t *__restrict__ pSrcBias, + const bool has_bias, float32_t *__restrict__ pDstC, uint32_t pad_top, + uint32_t pad_bottom, uint32_t pad_left, uint32_t pad_right, + float32_t *__restrict__ pContextBuffer, float32_t scale, int32_t zeroPoint); + #endif // __DEEPLOY_MATH_CONV_KERNEL_HEADER_ \ No newline at end of file diff --git a/TargetLibraries/PULPOpen/inc/kernel/Matmul.h b/TargetLibraries/PULPOpen/inc/kernel/Matmul.h index 9176801b..7ac421e4 100644 --- a/TargetLibraries/PULPOpen/inc/kernel/Matmul.h +++ b/TargetLibraries/PULPOpen/inc/kernel/Matmul.h @@ -14,4 +14,12 @@ void PULP_MatMul_fp32_fp32_fp32_unroll1x7(const float32_t *__restrict__ pSrcA, float32_t *__restrict__ pDstY, uint32_t M, uint32_t N, uint32_t O); +// Matmul against an int8 weight matrix with per-tensor affine dequantisation +// folded into the loop; see Matmul.c for why the scale factors out. +void PULP_MatMul_fp32_i8_fp32_unroll1x7(const float32_t *__restrict__ pSrcA, + const int8_t *__restrict__ pSrcB, + float32_t *__restrict__ pDstY, + uint32_t M, uint32_t N, uint32_t O, + float32_t scale, int32_t zeroPoint); + #endif // __DEEPLOY_MATH_MATMUL_KERNEL_HEADER_ \ No newline at end of file diff --git a/TargetLibraries/PULPOpen/inc/kernel/gemm.h b/TargetLibraries/PULPOpen/inc/kernel/gemm.h index 863bb210..42415241 100644 --- a/TargetLibraries/PULPOpen/inc/kernel/gemm.h +++ b/TargetLibraries/PULPOpen/inc/kernel/gemm.h @@ -14,6 +14,15 @@ void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA, const float32_t *__restrict__ pDstC, float32_t *__restrict__ pDstY, uint32_t M, uint32_t N, uint32_t O, uint32_t transA, - uint32_t transB); + uint32_t transB, uint32_t biasStride); + +// Gemm with an int8 weight, dequantised in-kernel; see Gemm.c. +void PULP_Gemm_fp32_i8_fp32_fp32(const float32_t *__restrict__ pSrcA, + const int8_t *__restrict__ pSrcB, + const float32_t *__restrict__ pDstC, + float32_t *__restrict__ pDstY, uint32_t M, + uint32_t N, uint32_t O, uint32_t transA, + uint32_t transB, uint32_t biasStride, + float32_t scale, int32_t zeroPoint); #endif // __DEEPLOY_MATH_GEMM_KERNEL_HEADER_ \ No newline at end of file diff --git a/TargetLibraries/PULPOpen/src/Convolution_fp32.c b/TargetLibraries/PULPOpen/src/Convolution_fp32.c index 5dee8a7c..14946160 100644 --- a/TargetLibraries/PULPOpen/src/Convolution_fp32.c +++ b/TargetLibraries/PULPOpen/src/Convolution_fp32.c @@ -293,3 +293,150 @@ void PULP_Conv2d_Im2Col_fp32_fp32_fp32_HWC( } } } + +// HWC im2col convolution against an int8 weight, dequantising inside the loop. +// +// Same reasoning as PULP_MatMul_fp32_i8_fp32_unroll1x7: a weight-only-quantised +// model would otherwise materialise the dequantised weight as its own tensor. +// For the tokenizer's conv that is 288 KB of fp32 for a 72 KB int8 constant, +// and profiling puts Conv's pre-kernel time at 143.90M cycles against 0.87M for +// the same model unquantised -- the kernel itself is unchanged at 46.20M +// vs 45.85M, so the entire difference is moving that weight. +// +// The per-tensor affine dequantisation factors out of the accumulation: +// +// sum_k a_k * (w_k - zp) * scale == scale * (sum_k a_k * w_k - zp * sum_k +// a_k) +// +// so it costs one multiply per output element. sum_k a_k is over the im2col +// window and is shared by every output channel, hence hoisted above the channel +// loop; with zeroPoint 0 it is skipped entirely. +void PULP_Conv2d_Im2Col_fp32_i8_fp32_HWC( + const float32_t *__restrict__ pSrcA, uint32_t H, uint32_t W, uint32_t C, + const int8_t *__restrict__ pSrcB, uint32_t F_total, uint32_t P, uint32_t Q, + uint32_t SP, uint32_t SQ, const float32_t *__restrict__ pSrcBias, + const bool has_bias, float32_t *__restrict__ pDstC, uint32_t pad_top, + uint32_t pad_bottom, uint32_t pad_left, uint32_t pad_right, + float32_t *__restrict__ pContextBuffer, float32_t scale, + int32_t zeroPoint) { + + // Compute core + int8_t core_id = pi_core_id(); + int8_t log2Core = LOG2(NUM_CORES); + + // Compute the chunk size for each core + uint16_t ch_out_chunk = + (F_total >> log2Core) + ((F_total & (NUM_CORES - 1)) != 0); + uint16_t ch_out_start = MIN(ch_out_chunk * core_id, F_total); + uint16_t ch_out_stop = MIN(ch_out_start + ch_out_chunk, F_total); + uint16_t ch_out_count = ch_out_stop - ch_out_start; + + if (ch_out_count == 0) { + return; + } + + // Pointer to the weights for the current core + const int8_t *weight_ptr = pSrcB + ch_out_start * C * P * Q; + + uint32_t im2col_size_per_core = C * P * Q; + float32_t *im2col_buffer = pContextBuffer + core_id * im2col_size_per_core; + + // Compute the output dimensions + uint32_t H_out = (H + pad_top + pad_bottom - P) / SP + 1; + uint32_t W_out = (W + pad_left + pad_right - Q) / SQ + 1; + uint32_t kernel_size = P * Q * C; + + // Compute the output + if (has_bias) { + for (uint32_t h_out = 0; h_out < H_out; h_out++) { + for (uint32_t w_out = 0; w_out < W_out; w_out++) { + int32_t h_in_start = h_out * SP - pad_top; + int32_t w_in_start = w_out * SQ - pad_left; + + for (uint32_t p = 0; p < P; p++) { + int32_t h_in = h_in_start + p; + + for (uint32_t q = 0; q < Q; q++) { + int32_t w_in = w_in_start + q; + + for (uint32_t c = 0; c < C; c++) { + if (h_in >= 0 && h_in < (int32_t)H && w_in >= 0 && + w_in < (int32_t)W) { + uint32_t in_idx = (h_in * W + w_in) * C + c; + im2col_buffer[p * Q * C + q * C + c] = pSrcA[in_idx]; + } else { + im2col_buffer[p * Q * C + q * C + c] = 0.0f; + } + } + } + } + + float32_t im2col_sum = 0.0f; + if (zeroPoint != 0) { + for (uint32_t k = 0; k < kernel_size; k++) { + im2col_sum += im2col_buffer[k]; + } + } + for (uint32_t f = ch_out_start; f < ch_out_stop; f++) { + float32_t sum = 0.0f; + const int8_t *local_weight_ptr = + weight_ptr + (f - ch_out_start) * kernel_size; + + for (uint32_t k = 0; k < kernel_size; k++) { + sum += im2col_buffer[k] * (float32_t)local_weight_ptr[k]; + } + + uint32_t out_idx = (h_out * W_out + w_out) * F_total + f; + + pDstC[out_idx] = + scale * (sum - (float32_t)zeroPoint * im2col_sum) + pSrcBias[f]; + } + } + } + } else { + for (uint32_t h_out = 0; h_out < H_out; h_out++) { + for (uint32_t w_out = 0; w_out < W_out; w_out++) { + int32_t h_in_start = h_out * SP - pad_top; + int32_t w_in_start = w_out * SQ - pad_left; + + for (uint32_t p = 0; p < P; p++) { + int32_t h_in = h_in_start + p; + + for (uint32_t q = 0; q < Q; q++) { + int32_t w_in = w_in_start + q; + + for (uint32_t c = 0; c < C; c++) { + if (h_in >= 0 && h_in < (int32_t)H && w_in >= 0 && + w_in < (int32_t)W) { + uint32_t in_idx = (h_in * W + w_in) * C + c; + im2col_buffer[p * Q * C + q * C + c] = pSrcA[in_idx]; + } else { + im2col_buffer[p * Q * C + q * C + c] = 0.0f; + } + } + } + } + + float32_t im2col_sum = 0.0f; + if (zeroPoint != 0) { + for (uint32_t k = 0; k < kernel_size; k++) { + im2col_sum += im2col_buffer[k]; + } + } + for (uint32_t f = ch_out_start; f < ch_out_stop; f++) { + float32_t sum = 0.0f; + const int8_t *local_weight_ptr = + weight_ptr + (f - ch_out_start) * kernel_size; + + for (uint32_t k = 0; k < kernel_size; k++) { + sum += im2col_buffer[k] * (float32_t)local_weight_ptr[k]; + } + + uint32_t out_idx = (h_out * W_out + w_out) * F_total + f; + + pDstC[out_idx] = scale * (sum - (float32_t)zeroPoint * im2col_sum); + } + } + } + } +} diff --git a/TargetLibraries/PULPOpen/src/Gemm.c b/TargetLibraries/PULPOpen/src/Gemm.c index a46f8ac6..c4631e85 100644 --- a/TargetLibraries/PULPOpen/src/Gemm.c +++ b/TargetLibraries/PULPOpen/src/Gemm.c @@ -7,12 +7,19 @@ #include "DeeployPULPMath.h" #include "pmsis.h" +// biasStride is the row stride of pDstC: O for a full [M,O] bias, and 0 for a +// broadcast [O] bias, where every output row adds the same vector. A +// transformer stores its Linear bias as [O] but GEMMLayer.computeShapes used to +// widen the C operand to [M,O], materialising the same O values M times -- 32 +// KB per 128-wide bias at 64 tokens, against 512 B of data. Passing a stride +// keeps one kernel for both layouts, and biasStride == O reproduces the +// previous behaviour exactly. void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA, const float32_t *__restrict__ pSrcB, const float32_t *__restrict__ pDstC, float32_t *__restrict__ pDstY, uint32_t M, uint32_t N, uint32_t O, uint32_t transA, - uint32_t transB) { + uint32_t transB, uint32_t biasStride) { int8_t core_id = pi_core_id(); int8_t log2Core = LOG2(NUM_CORES); @@ -35,7 +42,8 @@ void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA, for (uint32_t i = M_start; i < M_end; ++i) { const float32_t *__restrict__ a_row = &pSrcA[i * N]; float32_t *__restrict__ y_row = &pDstY[i * O]; - const float32_t *__restrict__ c_row = has_bias ? &pDstC[i * O] : NULL; + const float32_t *__restrict__ c_row = + has_bias ? &pDstC[i * biasStride] : NULL; uint32_t j = 0; @@ -85,7 +93,8 @@ void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA, for (uint32_t i = M_start; i < M_end; ++i) { float32_t *__restrict__ y_row = &pDstY[i * O]; - const float32_t *__restrict__ c_row = has_bias ? &pDstC[i * O] : NULL; + const float32_t *__restrict__ c_row = + has_bias ? &pDstC[i * biasStride] : NULL; uint32_t j = 0; for (; j < O_unroll; j += 6) { @@ -181,7 +190,8 @@ void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA, for (uint32_t i = M_start; i < M_end; ++i) { const float32_t *__restrict__ a_row = &pSrcA[i * N]; float32_t *__restrict__ y_row = &pDstY[i * O]; - const float32_t *__restrict__ c_row = has_bias ? &pDstC[i * O] : NULL; + const float32_t *__restrict__ c_row = + has_bias ? &pDstC[i * biasStride] : NULL; uint32_t j = 0; for (; j < O_unroll; j += 6) { @@ -266,7 +276,8 @@ void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA, for (uint32_t i = M_start; i < M_end; ++i) { float32_t *__restrict__ y_row = &pDstY[i * O]; - const float32_t *__restrict__ c_row = has_bias ? &pDstC[i * O] : NULL; + const float32_t *__restrict__ c_row = + has_bias ? &pDstC[i * biasStride] : NULL; uint32_t j = 0; for (; j < O_unroll; j += 6) { @@ -351,4 +362,54 @@ void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA, } } } -} \ No newline at end of file +} + +// Gemm against an int8 weight, dequantising inside the loop. Same factorisation +// as PULP_MatMul_fp32_i8_fp32_unroll1x7: the per-tensor affine dequantisation +// comes out of the accumulation, so it costs one multiply per output element. +// +// This exists because the backward pass reads the same weight as the forward +// one. If only the forward MatMul takes the int8 binding, the shared constant +// gets typed from whichever binding the backward Gemm selected -- and a fp32 +// binding there re-types the weight to float, quadrupling its allocation. Both +// consumers have to agree. +void PULP_Gemm_fp32_i8_fp32_fp32(const float32_t *__restrict__ pSrcA, + const int8_t *__restrict__ pSrcB, + const float32_t *__restrict__ pDstC, + float32_t *__restrict__ pDstY, uint32_t M, + uint32_t N, uint32_t O, uint32_t transA, + uint32_t transB, uint32_t biasStride, + float32_t scale, int32_t zeroPoint) { + int8_t core_id = pi_core_id(); + int8_t log2Core = LOG2(NUM_CORES); + uint32_t M_chunk = (M >> log2Core) + ((M & (NUM_CORES - 1)) != 0); + uint32_t M_start = MIN(core_id * M_chunk, M); + uint32_t M_end = MIN(M_start + M_chunk, M); + if (M_end <= M_start) { + return; + } + const uint32_t has_bias = (pDstC != NULL); + + for (uint32_t i = M_start; i < M_end; ++i) { + const float32_t *__restrict__ a_row = &pSrcA[i * N]; + float32_t *__restrict__ y_row = &pDstY[i * O]; + const float32_t *__restrict__ c_row = + has_bias ? &pDstC[i * biasStride] : NULL; + + float32_t sum_a = 0.0f; + if (zeroPoint != 0) { + for (uint32_t k = 0; k < N; ++k) { + sum_a += a_row[k]; + } + } + const float32_t correction = scale * (float32_t)zeroPoint * sum_a; + + for (uint32_t j = 0; j < O; ++j) { + float32_t sum = 0.0f; + for (uint32_t k = 0; k < N; ++k) { + sum += a_row[k] * (float32_t)pSrcB[k * O + j]; + } + y_row[j] = scale * sum - correction + (has_bias ? c_row[j] : 0.0f); + } + } +} diff --git a/TargetLibraries/PULPOpen/src/Matmul.c b/TargetLibraries/PULPOpen/src/Matmul.c index 639da75f..0c07e89a 100644 --- a/TargetLibraries/PULPOpen/src/Matmul.c +++ b/TargetLibraries/PULPOpen/src/Matmul.c @@ -82,4 +82,89 @@ void PULP_MatMul_fp32_fp32_fp32_unroll1x7(const float32_t *__restrict__ pSrcA, local_pDstY[i * O + j] = sum; } } -} \ No newline at end of file +} +// Matmul against an int8 weight matrix, dequantising inside the loop. +// +// A weight-only-quantised model stores the frozen weight as int8 and needs it +// as float for the fp32 matmul. Emitting that as a separate Dequant node +// materialises the whole dequantised matrix -- on CCT-QLoRA, 1088 KB of fp32 +// tensors that exist only to be consumed by the very next node, which then has +// to be kept alive or recomputed. QLoRA itself never does this: it dequantises +// inside the matmul kernel and discards the value immediately. +// +// The dequantisation is affine and per-tensor, so it factors out of the inner +// loop entirely: +// +// sum_k a_k * (b_k - zp) * scale == scale * (sum_k a_k * b_k - zp * sum_k +// a_k) +// +// leaving one multiply per output element instead of one per +// multiply-accumulate. With zeroPoint 0 the second term vanishes and the loop +// is the fp32 loop with an int8 load. Numerically identical to Dequant followed +// by matmul, up to the order of the same floating-point operations. +void PULP_MatMul_fp32_i8_fp32_unroll1x7(const float32_t *__restrict__ pSrcA, + const int8_t *__restrict__ pSrcB, + float32_t *__restrict__ pDstY, + uint32_t M, uint32_t N, uint32_t O, + float32_t scale, int32_t zeroPoint) { + int8_t core_id = pi_core_id(); + int8_t log2Core = LOG2(NUM_CORES); + uint32_t M_chunk = (M >> log2Core) + ((M & (NUM_CORES - 1)) != 0); + uint32_t M_start = MIN(core_id * M_chunk, M); + uint32_t M_end = MIN(M_start + M_chunk, M); + uint32_t M_size = M_end - M_start; + + if (M_size == 0) { + return; + } + + const float32_t *local_pSrcA = pSrcA + M_start * N; + float32_t *local_pDstY = pDstY + M_start * O; + uint32_t O_block = O - (O % 7); + + for (uint32_t i = 0; i < M_size; i++) { + // sum_a is only needed when the zero point is non-zero; the compiler drops + // it when zeroPoint is a compile-time 0 after inlining, and it costs one + // add per element otherwise. + float32_t sum_a = 0.0f; + if (zeroPoint != 0) { + for (uint32_t k = 0; k < N; k++) { + sum_a += local_pSrcA[i * N + k]; + } + } + const float32_t correction = scale * (float32_t)zeroPoint * sum_a; + + for (uint32_t j = 0; j < O_block; j += 7) { + float32_t sum0 = 0.0f, sum1 = 0.0f, sum2 = 0.0f, sum3 = 0.0f; + float32_t sum4 = 0.0f, sum5 = 0.0f, sum6 = 0.0f; + + for (uint32_t k = 0; k < N; k++) { + float32_t a0 = local_pSrcA[i * N + k]; + const int8_t *b_row = &pSrcB[k * O + j]; + sum0 += a0 * (float32_t)b_row[0]; + sum1 += a0 * (float32_t)b_row[1]; + sum2 += a0 * (float32_t)b_row[2]; + sum3 += a0 * (float32_t)b_row[3]; + sum4 += a0 * (float32_t)b_row[4]; + sum5 += a0 * (float32_t)b_row[5]; + sum6 += a0 * (float32_t)b_row[6]; + } + + local_pDstY[i * O + (j + 0)] = scale * sum0 - correction; + local_pDstY[i * O + (j + 1)] = scale * sum1 - correction; + local_pDstY[i * O + (j + 2)] = scale * sum2 - correction; + local_pDstY[i * O + (j + 3)] = scale * sum3 - correction; + local_pDstY[i * O + (j + 4)] = scale * sum4 - correction; + local_pDstY[i * O + (j + 5)] = scale * sum5 - correction; + local_pDstY[i * O + (j + 6)] = scale * sum6 - correction; + } + + for (uint32_t j = O_block; j < O; j++) { + float32_t sum = 0.0f; + for (uint32_t k = 0; k < N; k++) { + sum += local_pSrcA[i * N + k] * (float32_t)pSrcB[k * O + j]; + } + local_pDstY[i * O + j] = scale * sum - correction; + } + } +}