[fx/profiler] assigned UUID to each unrecorded tensor/ improved performance on GPT-2 #1679
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
super-dainiu
changed the title
super-dainiu
commented
Oct 5, 2022
Comment on lines
+51
to
63
| # TODO(super-dainiu): removed redundant items, currently all of them are necessary for development | ||
|
|
||
| fwd_flop: int = 0 | ||
| fwd_time: float = 0.0 | ||
| bwd_flop: int = 0 | ||
| bwd_time: float = 0.0 | ||
| save_fwd_in: bool = False | ||
| fwd_in: List = field(default_factory=list) | ||
| fwd_tmp: List = field(default_factory=list) | ||
| fwd_out: List = field(default_factory=list) | ||
| fwd_mem_tmp: int = 0 | ||
| fwd_mem_out: int = 0 | ||
| bwd_mem_tmp: int = 0 |
There was a problem hiding this comment.
Just for compatibility of different InfoProp
There was a problem hiding this comment.
Now all fwd of MetaInfoProp are saved in List
super-dainiu
commented
Oct 5, 2022
| # If there is an argument that this `call_function` is inplace, we should | ||
| # still run the profiling but discard some results regarding `target` | ||
| global do_not_cache |
There was a problem hiding this comment.
This happens if we replace inplace=True with inplace=False, because some tensors might be saved for backward during inplace=False, and it will be cached.
super-dainiu
commented
Oct 5, 2022
Comment on lines
+12
to
+15
| def set_uuid(x): | ||
| if isinstance(x, torch.Tensor): | ||
| if not hasattr(x, 'uuid'): | ||
| setattr(x, 'uuid', uuid.uuid4()) |
There was a problem hiding this comment.
uuid is set if and only if the tensor has no uuid
super-dainiu
commented
Oct 5, 2022
| @@ -15,7 +19,7 @@ | |||
| with_codegen = False | |||
|
|
|||
|
|
|||
| @pytest.mark.skip(reason='TODO: modify calculations in rotor') | |||
| @pytest.mark.skip(reason='TODO: modify the logger') | |||
super-dainiu
commented
Oct 5, 2022
| codegen = ActivationCheckpointCodeGen() | ||
| gm.graph.set_codegen(codegen) | ||
| if solver == solver_rotor: | ||
| gm = solver(gm, data, mem_limit=500 * 1024 * 1024, mem_slots=500) | ||
| gm = solver(gm, data, mem_limit=500 * 1024 * 1024, mem_slots=500, force_python=True) |
Cypher30
reviewed
Oct 6, 2022
| codegen = ActivationCheckpointCodeGen() | ||
| gm.graph.set_codegen(codegen) | ||
| if solver == solver_rotor: | ||
| gm = solver(gm, data, mem_limit=500 * 1024 * 1024, mem_slots=500) | ||
| gm = solver(gm, data, mem_limit=500 * 1024 * 1024, mem_slots=500, force_python=True) |
There was a problem hiding this comment.
Currently I cannot reproduce this problem.
super-dainiu
commented
Oct 6, 2022
Comment on lines
+1
to
+42
| # for PyTorch 1.11 compatibility uses | ||
| import torch | ||
| from torch.fx import Node, GraphModule | ||
| from typing import Union, Dict, List, Tuple | ||
|
|
||
| __all__ = ["calculate_fwd_in", "calculate_fwd_tmp", "calculate_fwd_out"] | ||
|
|
||
|
|
||
| def calculate_fwd_in(n: Node) -> bool: | ||
| """A helper function to calculate `fwd_in` | ||
|
|
||
| Args: | ||
| n (Node): a node from the graph | ||
|
|
||
| Returns: | ||
| save_fwd_in (bool): the result of `save_fwd_in` | ||
| """ | ||
| return n.meta['save_fwd_in'] | ||
|
|
||
|
|
||
| def calculate_fwd_tmp(n: Node) -> int: | ||
| """A helper function to calculate `fwd_tmp` | ||
|
|
||
| Args: | ||
| n (Node): a node from the graph | ||
|
|
||
| Returns: | ||
| fwd_tmp (int): the result of `fwd_tmp` | ||
| """ | ||
| return n.meta["fwd_mem_tmp"] | ||
|
|
||
|
|
||
| def calculate_fwd_out(n: Node) -> int: | ||
| """A helper function to calculate `fwd_out` | ||
|
|
||
| Args: | ||
| n (Node): a node from the graph | ||
|
|
||
| Returns: | ||
| fwd_out (int): the result of `fwd_out` | ||
| """ | ||
| return n.meta['fwd_mem_out'] |
There was a problem hiding this comment.
this is for compatibility uses because metainfoprop for torch 1.11 don't save tensors with uuid.
super-dainiu
commented
Oct 6, 2022
Comment on lines
+168
to
+180
| def _get_fwd_mem_tmp(node: List[Node]) -> int: | ||
| """Get the forward temp memory of a node | ||
| This could be done by subtracting the saved activation from all output of a node | ||
|
|
||
| Args: | ||
| node (List[Node]): List of torch.fx Node, | ||
| indicates a node in linearized graph | ||
|
|
||
| Returns: | ||
| int: forward temp memory, unit Byte | ||
| """ | ||
| n = node[-1] | ||
| return activation_size(n.meta['fwd_out']) - calculate_fwd_out(n) |
There was a problem hiding this comment.
Now fwd_out that are not saved for next node may be regarded as fwd_tmp
Cypher30
approved these changes
Oct 11, 2022
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What's new?
fwd_in,fwd_tmpandfwd_out. Hopefully, this will avoid a lot of misunderstandings.Bugs in existing code
C version results
Force Python results
What is wrong with this? ✔
Logger always requires a launch, some modifications should be made either on the logger or on
test_linearize.py. cc @Cypher30Tests
Passed all the tests/test_fx tests locally.
