mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
7a076b7b93
For a while now, we've been re-running our functionalization analysis pass twice - once for get metadata when dedup'ing, and an entire second time during aot_dispatch_base/autograd. This should also probably speed up compile times pretty noticeably, since we're going from: (a) inference-only trace case: 3 fw traces -> 2 fw traces (b) autograd trace case: 2 fw traces + 1 joint trace -> 1 fw trace + 1 joint trace Pull Request resolved: https://github.com/pytorch/pytorch/pull/95992 Approved by: https://github.com/ezyang
51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the BSD-style license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
"""
|
|
Global flags for aot autograd
|
|
"""
|
|
import os
|
|
import sys
|
|
import logging
|
|
|
|
use_functionalize = True
|
|
|
|
use_fake_tensor = True
|
|
|
|
# can be useful for debugging if we are incorrectly creating meta fake tensors
|
|
fake_tensor_allow_meta = os.environ.get("FAKE_ALLOW_META", True)
|
|
|
|
# Enables optional asserts in hotpath code to check for errors. If
|
|
# you are seeing weird accuracy problems, try turning this on.
|
|
# This is currently off by default as it will harm tracing time,
|
|
# but it is on by default for aot_eager.
|
|
debug_assert = False
|
|
|
|
debug_fake_cross_ref = os.environ.get("AOT_FAKE_CROSSREF", False)
|
|
|
|
debug_partitioner = os.environ.get("AOT_PARTITIONER_DEBUG", False)
|
|
# Prints out forward + backwards FX graphs
|
|
debug_graphs = os.environ.get("AOT_FX_GRAPHS", False)
|
|
# Prints out joint graph traced, before partitioning
|
|
debug_joint = os.environ.get("AOT_FX_GRAPHS_JOINT", False)
|
|
|
|
static_weight_shapes = True
|
|
|
|
# Applies CSE to the graph before partitioning
|
|
cse = True
|
|
|
|
# Restricts the amount of computation AOTAutograd can do.
|
|
max_dist_from_bw = 3
|
|
|
|
log_level = (
|
|
logging.DEBUG if debug_partitioner or debug_graphs or debug_joint else logging.INFO
|
|
)
|
|
|
|
from .._dynamo.config_utils import install_config_module
|
|
|
|
# adds patch, save_config, invalid config checks, etc
|
|
install_config_module(sys.modules[__name__])
|