Freezing Torchscript modules #32178

bzinodev · 2020-01-14T19:16:03Z

This patch enables folding GetAttr nodes with their corresponding
values. _jit_pass_freeze_module API returns a new TorchScipt module
where all function calls and get attributes are inlined.
Usage:

frozen_model = torch._C._freeze_module(scrited_model._c)
frozen_model.forward(...)

This API currently optimizes the forward method. We will follow up to
to preserve and optimize methods and attributes that are annotated as
@torch.jit.interface.

Several future improvements to JIT optimizations are required to maximize
clean up/de-sugar the graph and eliminate redundancies.
Ideally, we want to produce a graph that can easily be lowered to
GLOW and other low-level backends.
__

kostmo · 2020-01-14T19:18:54Z

💊 CircleCI build failures summary and remediations

As of commit 597d63d (more details on the Dr. CI page):

Commit 597d63d was recently pushed. Waiting for builds...

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker.

This comment has been revised 85 times.

facebook-github-bot

@bzinodev has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@bzinodev has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

test/test_jit.py

torch/csrc/jit/init.cpp

torch/csrc/jit/passes/freeze_module.cpp

eellison

Only took a brief look at alias analysis changes. As far as I understand, the semantics of freezing are:

"I will only run this frozen method after freezing" - which means that we must consider mutable attributes within the method, but any attributes which are not mutated we can inline.

If the pass correctly takes account of attributes which aren't mutated and then inlines them, why are there being changes made to alias analysis?

Edit: this still doesn't handle interface calls. I think we should just throw if we see one for now, since interfaces still aren't part of the public api.

torch/csrc/jit/passes/alias_analysis.cpp

suo

Looking pretty close. I left some commentary inline. We should also do a quick style review after this; but I don't want to clutter up the discussion on the actual logic.

torch/csrc/jit/passes/alias_analysis.cpp

torch/csrc/jit/passes/freeze_module.cpp

test/jit/test_module_interface.py

eellison

I think in order to inline attributes of a module the following needs to happen.

For a module value %ModVal of type Mod with a Tensor field weights:

no value with Mod type has any prim::SetAttr[field="weights"] (*1)
all prim::GetAttr["weights"](%ModVal) can be resolved statically
- For example, if %ModVal is an output of a control flow node we cannot resolve all GetAttr
weights is not aliased by any other Tensor contained within the frozen module, including in Lists and Tuples of Tensors (*2)
- You can check this by checking the Tensor storage of each tensor
at this point, you can replace all prim::GetAttr[field="weights"] with a single top-level prim::FrozenGetAttr. This preserves aliasing. If the vallue doesn't have any writers it can be inlined as a constant.

(*1) You can relax this condition for a specific value of type MyMod if you can prove it doesn't alias any MyMod values that reassign weights, maybe better as a follow

(*2) You can relax this condition if you can prove the other aliases aren't mutated, also maybe better as a follow up.

Edit: the above correctness condition ensures that there is one single alias for an attribute Tensor, meaning that it will work even if the Tensor is mutated.

A separate correctness condition is to check that all aliases of an attribute Tensor are not mutated

torch/csrc/jit/passes/alias_analysis.cpp

Summary: Currently we compile constants directly into the graph, but when we fold ConvBn, we might need to change a constant bias to some Tensor, and this is not possible with bias being in the `__constants__` list, so we need to remove `bias` from the list. This might result in some minor perf regression but it should be unnoticeable and we can restore the perf after freeze feature is enabled for production models: #32178 Test Plan: python test/test_jit.py Reviewers: suo, mvz Subscribers: Tasks: Tags: [ghstack-poisoned]

test/test_jit.py

eellison

Just did a quick skim. While you are right that when we script a module we construct a new list when we convert a python list to an ivalue. However that doesn't mean that each list is a unique alias.

class MyMod(nn.Module):
     def __init__(self):
          self.x = [1, 2, 3]
          self.y = [1]

     def make_alias(self):
          self.x = self.y
...
mod = torch.jit.script(MyMod())
mod.make_alias()
freeze(mod)

Agreed this is a pretty unusual case. however the more important thing here is that aliasing of Tensors is preserved. So two lists can contain the same tensor.

I think we need to work on / finish up the hashing of potentially aliasing ivalues still.

torch/csrc/jit/passes/alias_analysis.cpp

aten/src/ATen/core/ivalue.h

test/jit/test_freezing.py

torch/csrc/jit/passes/freeze_module.cpp

facebook-github-bot

@bzinodev has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

Summary: Currently we compile constants directly into the graph, but when we fold ConvBn, we might need to change a constant bias to some Tensor, and this is not possible with bias being in the `__constants__` list, so we need to remove `bias` from the list. This might result in some minor perf regression but it should be unnoticeable and we can restore the perf after freeze feature is enabled for production models: pytorch/pytorch#32178 Test Plan: python test/test_jit.py Reviewers: suo, mvz Subscribers: Tasks: Tags: ghstack-source-id: 8cdd38a Pull Request resolved: pytorch/pytorch#32543

facebook-github-bot

@bzinodev has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@bzinodev has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

eellison

🚢🚢🚢🚢🚢🚢🚢🚢🚢

Looks great, let's ship this! I left some comments that I would like addressed / responded to before landing, but let's merge this. It covers all of the reasonable edge cases, and as you have run into, requires some structural changes to Alias Analysis to handle the remaining degenerate cases that shouldn't block the PR, + we safely error out in those cases anyway.

Maybe you can help me here, but let's at least establish what some of the follows up:

maybe clean up / move public apis on ivalue into freezing
once alias analysis has fine-grained containment tracking, we can improve freezing here instead of erroring out
remove methods, attributes

anything else ?

eellison · 2020-02-28T21:10:38Z

aten/src/ATen/core/ivalue.cpp

+      break;
+    case Tag::Future:
+    case Tag::Device:
+    case Tag::Object:


You can do object like this:

case Tag::Object: { auto obj_type = type()->expect<ClassType>(); auto obj_value = toObject(); auto attribute_names = obj_type->attributeNames(); for (const auto& name: attribute_names) { auto attribute = obj_value->getAttr(name); attribute.getSubValues(subValues); } }

correct! I wanted to be conservative. At this point freezing does not need this case. I will add in the future with a nice testcase.

eellison · 2020-02-28T21:15:10Z

aten/src/ATen/core/ivalue.cpp

+    case Tag::Object:
+    case Tag::PyObject:
+    case Tag::Uninitialized:
+    case Tag::Capsule:


You can remove Device and Uninitialized from this list. Uninitialized only exists as an IR construct, should never be an attribute, Device is immutable.

eellison · 2020-02-28T21:15:31Z

aten/src/ATen/core/ivalue.cpp

+          false, "sub ivalue is nat enabled for: ", this->tagKind());
+      // Fall through
+    default:
+      // don't record scalars.


nit: /don't record scalars/ don't record immutable types

eellison · 2020-02-28T21:34:43Z

aten/src/ATen/core/ivalue.h

+  struct HashIValue {
+    size_t operator()(const IValue& val) const {
+      if (val.isTensor()) {
+        return 0;


If you look at the storage is_alias_of implementation, all it does is check the storage pointer. https://github.com/pytorch/pytorch/blob/master/c10/core/Storage.h#L152 We can do that here too.

#include <torch/csrc/utils/hash.h> if (val.isTensor()) { return reinterpret_cast<size_t>(val.toTensor().storage().unsafeGetStorageImpl()); } ...

paging @smessmer that i did this right

eellison · 2020-02-28T21:37:26Z

aten/src/ATen/core/ivalue.h

    return payload.as_intrusive_ptr;
  }

  TypePtr type() const;

+  size_t hash() const {
+    return payload.as_int;


This is hashing for a very specific question - do these ivalues alias. You could imagine another user of this api to be asking, "are these two tensors the same", and they would be using this wrong. I would prefer if we just moved this all into the freezing files instead of adding APIs to ivalue.

I am renaming the HashIValue to HashAliasedIValue and CompIValue to CompAliasedIValue. Hope this clear enough :)

eellison · 2020-02-28T22:34:47Z

torch/csrc/jit/passes/freeze_module.cpp

+          TORCH_INTERNAL_ASSERT(attrModule.hasattr(name));
+          Value* paramConst = nullptr;
+          auto I = attrValues.find(attrModule._ivalue());
+          if (I != attrValues.end()) {


nit: the I / II variable names are confusing to read.

Agreed, sorry in LLVM we capitalize variables. I am renaming I to iter and II to iter2

eellison · 2020-02-28T22:42:28Z

torch/csrc/jit/passes/freeze_module.cpp

+  script::Module& module_;
+
+  // Contains the attributes names (e.g. {"self", "subModule", "a"}
+  std::deque<std::string> names_;


Not sure why this is a class attribute, probably shouldn't be

I used it to pretty print (e.g. %self.sub.conv = prim::constant(...). the time of computing the name and inserting the new node are far apart.

eellison · 2020-02-28T22:43:42Z

torch/csrc/jit/passes/freeze_module.cpp

+    }
+  }
+
+  IValue overrideGradient(IValue attr) {


Can we not just call getSubValues(), and iterate through tensor subvalues, in-place removing gradient ?

Yes cleaner and it handles missing cases.

Eliminate tensor detach?

eellison · 2020-02-28T22:44:17Z

test/jit/test_freezing.py

+                self.c = (self.a, 10)
+
+            def forward(self, x):
+                self.b[1] += 10


eellison · 2020-02-28T22:48:02Z

torch/csrc/jit/passes/freeze_module.cpp

+  void run(std::shared_ptr<Graph>& graph) {
+    Inline(*graph);
+    propagateAttributes(graph);
+    runOptimization(graph, /* unroll? */ false);


I know quantization and potentially other parties wanted to run optimizations immediately after freezing. It might be worth removing the runOptimization call here to make freezing more composable.

The optimization is important for the clean up. Getting rid of all unused attributes. I think it is better to return a clean graph and user can call whatever followup otimization post freezing.

This patch enables folding GetAttr nodes with their corresponding values. _jit_pass_freeze_module API returns a new TorchScipt module where all function calls and get attributes are inlined. Usage: frozen_model = torch._C._freeze_module(scrited_model._c) frozen_model.forward(...) This API currently optimizes the forward method. We will follow up to to preserve and optimize methods and attributes on demand. Several future improvements to JIT optimizations are required to maximize clean up/de-sugar the graph and eliminate redundancies. this is an important step toward producing a graph that can easily be lowered to GLOW and other low-level backends. __

This patch adds two APIs to delete methods from modules and module's types. NOTE these APIs is only for internal only uses and should be used only by freezing where a new module and type are being created.

facebook-github-bot

@bzinodev has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

eellison · 2020-03-03T22:39:12Z

test/jit/test_freezing.py

+        v.append(4)
+        m_s.a = v
+        m_s.eval()
+        m_f = torch._C._freeze_module(m_s._c)


Also, before this pass is made public we need a workflow that does not involve accessing private members (torch._C and m_s._c). This pass should be in torch.jit.script with proper documentation (e.g. https://github.com/pytorch/pytorch/blob/master/torch/jit/__init__.py#L1117).

eellison · 2020-03-03T22:46:05Z

torch/csrc/jit/passes/freeze_module.cpp

+  // folded.
+  // TODO: Determine if freezing in training mode is useful and further clarify
+  // its semantics.
+  TORCH_CHECK(!module.is_training());


Could we improve the error message here ?

Summary: This patch enables folding GetAttr nodes with their corresponding values. _jit_pass_freeze_module API returns a new TorchScipt module where all function calls and get attributes are inlined. Usage: frozen_model = torch._C._freeze_module(scrited_model._c) frozen_model.forward(...) This API currently optimizes the forward method. We will follow up to to preserve and optimize methods and attributes that are annotated as torch.jit.interface. Several future improvements to JIT optimizations are required to maximize clean up/de-sugar the graph and eliminate redundancies. Ideally, we want to produce a graph that can easily be lowered to GLOW and other low-level backends. __ Pull Request resolved: pytorch#32178 Differential Revision: D19419640 Pulled By: bzinodev fbshipit-source-id: 52baffaba9bca2cd60a8e747baa68d57711ad42b

bzinodev requested a review from apaszke as a code owner January 14, 2020 19:16

facebook-github-bot added the oncall: jit Add this issue/PR to JIT oncall triage queue label Jan 14, 2020

bzinodev force-pushed the freeze_module branch 3 times, most recently from 4c584d3 to beb2e40 Compare January 15, 2020 22:43

facebook-github-bot reviewed Jan 15, 2020

View reviewed changes

bzinodev requested review from ZolotukhinM, suo, zdevito, jerryzh168 and jackm321 January 15, 2020 22:45

bzinodev force-pushed the freeze_module branch 2 times, most recently from 35265b1 to 110a426 Compare January 15, 2020 23:13

bzinodev requested a review from resistor January 15, 2020 23:18

facebook-github-bot reviewed Jan 15, 2020

View reviewed changes

bzinodev force-pushed the freeze_module branch from 110a426 to bcf1ec7 Compare January 16, 2020 00:49

jerryzh168 reviewed Jan 16, 2020

View reviewed changes

test/test_jit.py Outdated Show resolved Hide resolved

jerryzh168 reviewed Jan 16, 2020

View reviewed changes

torch/csrc/jit/init.cpp Outdated Show resolved Hide resolved

jerryzh168 reviewed Jan 16, 2020

View reviewed changes

torch/csrc/jit/passes/freeze_module.cpp Show resolved Hide resolved

jerryzh168 reviewed Jan 16, 2020

View reviewed changes

torch/csrc/jit/passes/freeze_module.cpp Show resolved Hide resolved

bzinodev force-pushed the freeze_module branch from bcf1ec7 to 602a1b9 Compare January 16, 2020 03:56

bzinodev requested a review from eellison January 16, 2020 05:26

eellison reviewed Jan 16, 2020

View reviewed changes

torch/csrc/jit/passes/alias_analysis.cpp Outdated Show resolved Hide resolved

bzinodev force-pushed the freeze_module branch from 602a1b9 to 49e2be1 Compare January 17, 2020 01:50

suo requested changes Jan 17, 2020

View reviewed changes

suo reviewed Jan 21, 2020

View reviewed changes

torch/csrc/jit/passes/freeze_module.cpp Outdated Show resolved Hide resolved

torch/csrc/jit/passes/freeze_module.cpp Outdated Show resolved Hide resolved

test/jit/test_module_interface.py Show resolved Hide resolved

eellison reviewed Jan 22, 2020

View reviewed changes

torch/csrc/jit/passes/alias_analysis.cpp Outdated Show resolved Hide resolved

jerryzh168 mentioned this pull request Jan 23, 2020

[WIP][quant][graphmode] ConvBn Folding works with bias=None #32543

Closed

iseeyuan reviewed Feb 7, 2020

View reviewed changes

test/test_jit.py Outdated Show resolved Hide resolved

bzinodev force-pushed the freeze_module branch 2 times, most recently from 3242e64 to 5841e2a Compare February 12, 2020 22:07

eellison reviewed Feb 12, 2020

View reviewed changes

torch/csrc/jit/passes/alias_analysis.cpp Outdated Show resolved Hide resolved

aten/src/ATen/core/ivalue.h Show resolved Hide resolved

test/jit/test_freezing.py Outdated Show resolved Hide resolved

bzinodev force-pushed the freeze_module branch from 5841e2a to d171576 Compare February 12, 2020 23:20

ZolotukhinM reviewed Feb 12, 2020

View reviewed changes

torch/csrc/jit/passes/freeze_module.cpp Outdated Show resolved Hide resolved

bzinodev mentioned this pull request Feb 13, 2020

Freezing TorchScript Modules #33258

Closed

bzinodev linked an issue Feb 13, 2020 that may be closed by this pull request

Freezing TorchScript Modules #33258

Closed

bzinodev force-pushed the freeze_module branch from d171576 to 27e0c35 Compare February 13, 2020 01:00

facebook-github-bot reviewed Feb 13, 2020

View reviewed changes

bzinodev force-pushed the freeze_module branch 2 times, most recently from f2194d6 to 6a86724 Compare February 26, 2020 07:48

facebook-github-bot reviewed Feb 26, 2020

View reviewed changes

bzinodev force-pushed the freeze_module branch 2 times, most recently from cc8848c to be52a4d Compare February 28, 2020 19:48

eellison approved these changes Feb 28, 2020

View reviewed changes

Zino Benaissa added 3 commits February 29, 2020 10:00

Clean up frozen TorchScript Modules

7e8bd31

This patch adds two APIs to delete methods from modules and module's types. NOTE these APIs is only for internal only uses and should be used only by freezing where a new module and type are being created.

t

597d63d

bzinodev force-pushed the freeze_module branch from be52a4d to 597d63d Compare February 29, 2020 18:23

facebook-github-bot reviewed Mar 1, 2020

View reviewed changes

facebook-github-bot closed this in cab8772 Mar 2, 2020

eellison reviewed Mar 3, 2020

View reviewed changes

facebook-github-bot deleted the freeze_module branch July 13, 2020 17:55

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Freezing Torchscript modules #32178

Freezing Torchscript modules #32178

bzinodev commented Jan 14, 2020 •

edited

Loading

kostmo commented Jan 14, 2020 •

edited by dr-ci bot

Loading

facebook-github-bot left a comment

facebook-github-bot left a comment

eellison left a comment •

edited

Loading

suo left a comment

eellison left a comment •

edited

Loading

eellison left a comment

facebook-github-bot left a comment

facebook-github-bot left a comment

facebook-github-bot left a comment

eellison left a comment

eellison Feb 28, 2020

bzinodev Feb 29, 2020

eellison Feb 28, 2020

eellison Feb 28, 2020

eellison Feb 28, 2020

eellison Feb 28, 2020

bzinodev Mar 11, 2020

eellison Feb 28, 2020

bzinodev Mar 11, 2020

eellison Feb 28, 2020

bzinodev Mar 11, 2020

eellison Feb 28, 2020

bzinodev Mar 11, 2020

bzinodev Mar 11, 2020

eellison Feb 28, 2020

eellison Feb 28, 2020

bzinodev Mar 11, 2020

facebook-github-bot left a comment

eellison Mar 3, 2020

eellison Mar 3, 2020

Freezing Torchscript modules #32178

Freezing Torchscript modules #32178

Conversation

bzinodev commented Jan 14, 2020 • edited Loading

kostmo commented Jan 14, 2020 • edited by dr-ci bot Loading

💊 CircleCI build failures summary and remediations

facebook-github-bot left a comment

Choose a reason for hiding this comment

facebook-github-bot left a comment

Choose a reason for hiding this comment

eellison left a comment • edited Loading

Choose a reason for hiding this comment

suo left a comment

Choose a reason for hiding this comment

eellison left a comment • edited Loading

Choose a reason for hiding this comment

eellison left a comment

Choose a reason for hiding this comment

facebook-github-bot left a comment

Choose a reason for hiding this comment

facebook-github-bot left a comment

Choose a reason for hiding this comment

facebook-github-bot left a comment

Choose a reason for hiding this comment

eellison left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

facebook-github-bot left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bzinodev commented Jan 14, 2020 •

edited

Loading

kostmo commented Jan 14, 2020 •

edited by dr-ci bot

Loading

eellison left a comment •

edited

Loading

eellison left a comment •

edited

Loading