changes

Author: michaelbenayoun

optimum/graphcore/fx/utils.py

@@ -25,7 +25,7 @@
     MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES,
     MODEL_FOR_CTC_MAPPING_NAMES,
 )
-from transformers.utils.fx import HFTracer, get_concrete_args
+from transformers.utils.fx import HFAttribute, HFProxy, HFTracer, get_concrete_args
 
 from ..modeling_utils import PipelineMixin
 
@@ -34,11 +34,30 @@
     from transformers import PreTrainedModel
 
 
+# TODO: keep this until transformers >= 4.23.2
+class GCProxy(HFProxy):
+
+    @property
+    def dtype(self):
+        return self.__getattr__("dtype")
+
+    def __getattr__(self, k):
+        if k == "_metadata":
+            return self.__getattribute__(k)
+        # note: not added to the graph yet, if this is a method call
+        # we peephole optimize to the method invocation
+        hf_attribute = HFAttribute(self, k)
+        if hasattr(self, "_metadata"):
+            hf_attribute.install_metadata(getattr(self._metadata, k))
+        return hf_attribute
+
+
 class PipelinedTracer(HFTracer):
     # TODO: keep this until transformers >= 4.23.2
     _TORCH_METHODS_TO_PATCH = list(HFTracer._TORCH_METHODS_TO_PATCH)
     _TORCH_METHODS_TO_PATCH.append("clamp")
     _TORCH_METHODS_TO_PATCH.append("rand")
+    _TORCH_METHODS_TO_PATCH.append("finfo")
     """
     Tracer that enables tracing and transforming models to run them on IPUs.
     Compared to the HFTracer, this one adds the following features:
@@ -79,8 +98,9 @@ def proxy(self, node):
         # it is easier to use this one, and equivalent.
         node.parent_module_qualified_name = self.current_module_qualified_name[-1]
         node.parent_module_type = self.current_module_type[-1]
-        proxy = super().proxy(node)
-        return proxy
+        return GCProxy(node, self)
+        # return gc_proxy
+        return super().proxy(node)
 
     def call_module(self, m, forward, args, kwargs):
         # Could be done in a "cleaner" fashion by inlining the content of Tracer.call_module.
@@ -98,22 +118,22 @@ def call_module(self, m, forward, args, kwargs):
         return proxy
 
     def create_proxy(self, kind, target, args, kwargs, name=None, type_expr=None, proxy_factory_fn=None):
-        if self.root_is_in_half_precision:
-            float32_dtype_in_args = any(a is torch.float32 for a in args)
-            float32_dtype_in_kwargs = kwargs.get("dtype", None) is torch.float32
-            node_types_to_inspect = [
-                ("call_method", "to"),
-                ("call_function", torch.full),
-            ]
-            torch_methods_to_patched_version = {
-                orig: wrapped for (orig, wrapped) in self.patched_torch_methods.values()
-            }
-            for (k, t) in node_types_to_inspect:
-                if kind == k and target == torch_methods_to_patched_version.get(t, t):
-                    if float32_dtype_in_args:
-                        args = tuple(a if a is not torch.float32 else torch.float16 for a in args)
-                    if float32_dtype_in_kwargs:
-                        kwargs["dtype"] = torch.float16
+        # if self.root_is_in_half_precision:
+        #     float32_dtype_in_args = any(a is torch.float32 for a in args)
+        #     float32_dtype_in_kwargs = kwargs.get("dtype", None) is torch.float32
+        #     node_types_to_inspect = [
+        #         ("call_method", "to"),
+        #         ("call_function", torch.full),
+        #     ]
+        #     torch_methods_to_patched_version = {
+        #         orig: wrapped for (orig, wrapped) in self.patched_torch_methods.values()
+        #     }
+        #     for (k, t) in node_types_to_inspect:
+        #         if kind == k and target == torch_methods_to_patched_version.get(t, t):
+        #             if float32_dtype_in_args:
+        #                 args = tuple(a if a is not torch.float32 else torch.float16 for a in args)
+        #             if float32_dtype_in_kwargs:
+        #                 kwargs["dtype"] = torch.float16
         return super().create_proxy(
             kind, target, args, kwargs, name=name, type_expr=type_expr, proxy_factory_fn=proxy_factory_fn
         )
@@ -149,7 +169,6 @@ def symbolic_trace_with_pipelined_tracer(
     model: PipelineMixin,
     input_names: Optional[List[str]] = None,
 ) -> torch.fx.GraphModule:
-
     """
     Performs symbolic tracing on the model.
 

optimum/graphcore/models/deberta/modeling_deberta.py

@@ -37,7 +37,7 @@
 )
 from transformers.utils.fx import _gen_constructor_wrapper
 
-from ....fx.optimization import MergeLinears, compose
+from ....fx.optimization import MergeLinears, ReversibleTransformation, compose
 from ....utils import logging
 from ...fx import (
     DEFAULT_TRANSFORMATION_MANAGER,
@@ -46,6 +46,8 @@
     OutlineAttribute,
     RecomputationCheckpoint,
     VocabEmbeddingToSerializedEmbedding,
+    LinearToSerializedLinear,
+    TieWeights,
     symbolic_trace_pipelined_model,
 )
 from ...modeling_utils import OnehotGather, PipelineMixin, get_layer_ipu, register
@@ -107,7 +109,32 @@ def _get_rel_embedding(self):
     return self.rel_embeddings.weight + 0.0 if self.relative_attention else None
 
 
-gather_last_dim = FastGatherLastDim()
+def faster_gather_last_dim(input, dim, index, *args, **kwargs):
+    target = torch.zeros_like(index).to(input.dtype)
+    target.requires_grad_()
+    o = poptorch.custom_op(
+        [input, index],
+        "FastGatherLastDim",
+        "poptorch.custom_ops",
+        1,
+        example_outputs=[target],
+        attributes={"axis": -1},
+    )
+    return o[0]
+
+
+class ChangeTorchGather(ReversibleTransformation):
+    def transform(self, graph_module):
+        for node in graph_module.graph.nodes:
+            if node.op == "call_function" and node.target is torch.gather:
+                node.target = faster_gather_last_dim
+        return graph_module
+
+    def reverse(self, graph_module):
+        for node in graph_module.graph.nodes:
+            if node.op == "call_function" and node.target is faster_gather_last_dim:
+                node.target = torch.gather
+        return graph_module
 
 
 class IPUDisentangledSelfAttention(DisentangledSelfAttention):
@@ -124,8 +151,6 @@ class IPUDisentangledSelfAttention(DisentangledSelfAttention):
     def __init__(self, config):
         super().__init__(config)
         self.xsoftmax = XSoftmax(-1)
-        # self.gather_last_dim = FastGatherLastDim()
-        self.gather_last_dim = gather_last_dim
 
     def forward(
         self,
@@ -248,7 +273,8 @@ def disentangled_att_bias(self, query_layer, key_layer, relative_pos, rel_embedd
             index = c2p_pos.expand(
                 [query_layer.size(0), query_layer.size(1), query_layer.size(2), relative_pos.size(-1)]
             )
-            c2p_att = self.gather_last_dim(c2p_att, index)
+            # c2p_att = gather_last_dim(c2p_att, index)
+            c2p_att = torch.gather(c2p_att, -1, index)
             score += c2p_att
 
         # position->content
@@ -263,12 +289,12 @@ def disentangled_att_bias(self, query_layer, key_layer, relative_pos, rel_embedd
             p2c_pos = torch.clamp(-r_pos + att_span, 0, att_span * 2 - 1)
             index = p2c_pos.expand([query_layer.size(0), query_layer.size(1), key_layer.size(-2), key_layer.size(-2)])
             p2c_att = torch.matmul(key_layer, pos_query_layer.transpose(-1, -2))
-            p2c_att = self.gather_last_dim(p2c_att, index).transpose(-1, -2)
+            p2c_att = torch.gather(p2c_att, -1, index).transpose(-1, -2)
 
             if query_layer.size(-2) != key_layer.size(-2):
                 pos_index = relative_pos[:, :, :, 0].unsqueeze(-1)
                 index = pos_index.expand(pos_index, p2c_att, key_layer)
-                p2c_att = self.gather_last_dim(p2c_att, index)
+                p2c_att = torch.gather(p2c_att, -1, index)
             score += p2c_att
 
         return score
@@ -283,7 +309,6 @@ def change_modules_for_ipu(self, restore: bool):
                     del mod.xsoftmax
                 else:
                     mod.add_module("xsoftmax", XSoftmax(-1))
-                    mod.add_module("gather_last_dim", FastGatherLastDim())
             if restore:
                 if isinstance(mod, nn.Dropout):
                     mod.__class__ = StableDropout
@@ -302,10 +327,10 @@ def change_modules_for_ipu(self, restore: bool):
     def get_transformations(self):
         log_insertions = self.ipu_config.log_insertions
         layer_ipu = get_layer_ipu(self.ipu_config.layers_per_ipu)
+        # TODO: handle DebertaForMaskedLM
         transformations = [
             AddPoptorchBlock("Embedding", 0, "deberta.embeddings", log_insertions=log_insertions),
             OutlineAttribute("deberta.embeddings.LayerNorm", "Embedding"),
-            AddPoptorchBlock("Before Encoder", 0, "deberta.encoder", log_insertions=log_insertions),
             AddPoptorchBlocksInSeries(
                 "Encoder", layer_ipu, r"deberta.encoder.layer.[0-9]+", log_insertions=log_insertions
             ),
@@ -322,7 +347,13 @@ def get_transformations(self):
                 )
             )
         if self.ipu_config.embedding_serialization_factor > 1:
-            transformations.append(VocabEmbeddingToSerializedEmbedding())
+            if isinstance(self, DebertaForMaskedLM):
+                transformations += [
+                    LinearToSerializedLinear("cls.predictions.decoder"),
+                    TieWeights("deberta.embeddings.word_embeddings", "cls.predictions.decoder"),
+                ]
+            else:
+                transformations.append(VocabEmbeddingToSerializedEmbedding())
         return transformations
 
     def parallelize(self):
@@ -339,6 +370,7 @@ def parallelize(self):
         torch.nn.functional.one_hot = orig
         transformations = self.get_transformations()
         transformations += TRANSFORMATION_MANAGER.get_reversible_transformations(self.ipu_config.optimization_level)
+        transformations.append(ChangeTorchGather())
         composition = compose(*transformations)
         non_reversible_composition = TRANSFORMATION_MANAGER.compose_non_reversible_transformations(
             self.ipu_config.optimization_level

optimum/graphcore/models/gpt2/modeling_gpt2.py

@@ -131,7 +131,7 @@ def deparallelize(self):
 
 
 @register(GPT2LMHeadModel)
-class PipelinedGPT2LMHeadModel(GPT2LMHeadModel, GPT2PipelineMixin):
+class PipelinedGPT2LMHeadModel(GPT2PipelineMixin, GPT2LMHeadModel):
     def get_transformations(self):
         log_insertions = self.ipu_config.log_insertions
         layer_ipu = get_layer_ipu(self.ipu_config.layers_per_ipu)
@@ -241,7 +241,7 @@ def forward(
 
 
 @register(GPT2ForSequenceClassification)
-class PipelinedGPT2ForSequenceClassification(GPT2ForSequenceClassification, GPT2PipelineMixin):
+class PipelinedGPT2ForSequenceClassification(GPT2PipelineMixin, GPT2ForSequenceClassification):
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,
@@ -290,5 +290,5 @@ def forward(
 
 
 @register(GPT2ForTokenClassification)
-class PipelinedGPT2ForTokenClassification(GPT2ForTokenClassification, GPT2PipelineMixin):
+class PipelinedGPT2ForTokenClassification(GPT2PipelineMixin, GPT2ForTokenClassification):
     pass

optimum/graphcore/trainer.py

@@ -282,6 +282,8 @@ def __init__(
         if args.ipu_config_overrides:
             logger.info(f"Overriding IPU config: {args.ipu_config_overrides}")
             self.ipu_config.update_from_string(args.ipu_config_overrides)
+        if self.args.gradient_accumulation_steps is None:
+            self.args.gradient_accumulation_steps = self.ipu_config.gradient_accumulation_steps
         self.ipu_config.seed = self.args.seed
         self.opts = self.ipu_config.to_options(compile_only=args.compile_only)
         self.eval_opts = self.ipu_config.to_options(for_inference=True, compile_only=args.compile_only)
@@ -1116,7 +1118,7 @@ def _inner_training_loop(
         logger.info(f"  Num Epochs = {num_train_epochs}")
         logger.info(f"  Instantaneous batch size per device = {batch_size}")
         logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
-        logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+        logger.info(f"  Gradient Accumulation steps = {self.ipu_config.gradient_accumulation_steps}")
         logger.info(f"  Total optimization steps = {max_steps}")
 
         self.state.epoch = 0
@@ -1208,7 +1210,7 @@ def _inner_training_loop(
             steps_in_epoch = (
                 len(epoch_iterator)
                 if has_length(train_dataloader)
-                else args.max_steps * args.gradient_accumulation_steps
+                else args.max_steps * self.ipu_config.gradient_accumulation_steps
             )
 
             self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)

optimum/graphcore/training_args.py

@@ -750,8 +750,6 @@ def __post_init__(self):
         override_str = []
         if self.gradient_accumulation_steps is not None:
             override_str.append(f"gradient_accumulation_steps={self.gradient_accumulation_steps}")
-        else:
-            self.gradient_accumulation_steps = 1
 
         if self.auto_loss_scaling:
             override_str.append(f"auto_loss_scaling={self.auto_loss_scaling}")