Skip to content

Commit b34b192

Browse files
qihqipytorchmergebot
qihqi
authored and
pytorchmergebot
committed
Reland "Make debug_pkl smaller by only emitting unique traces." (pytorch#73368)
Summary: ## Original commit message: Pull Request resolved: pytorch#73368 debug_pkl file inside of pytorch's .pt file consists of a list of SourceRanges. Each SourceRange points to a Source which is a stack track, filename, and start, end numbers. Those are emitted in debug_pkl file as strings. Since many SourceRange shares the same source, the string for trace can be deduped. The newer format saves a set of unique traces in a tuple, then each SourceRange will save the offset of it's trace w.r.t. position in that tuple. (i.e. manually applying dictionary compression). The above helps with smaller file size. On loading, if we copy each trace to Source as string the runtime memory would still blowup. To mitigate this, we use SourceView directly instead of source which will take the reference of string inside of Deserializer and make that into string_view. This is safe because Deserializer is hold by Unpickler by shared_ptr, and Unpickler is also hold by shared_ptr by another Source object. That Source object will be alive during the model construction. Test Plan: ## Original Test plan unit test Took original file (312271638_930.predictor.disagg.local); loaded with `torch.jit.load` save again with `torch.jit.save`. Unzip both, look at contents: ``` [qihan@devvm5585.vll0 ~]$ du archive -h 4.0K archive/xl_model_weights 3.7M archive/extra 8.0K archive/code/__torch__/caffe2/torch/fb/model_transform/splitting 8.0K archive/code/__torch__/caffe2/torch/fb/model_transform 8.0K archive/code/__torch__/caffe2/torch/fb 8.0K archive/code/__torch__/caffe2/torch 8.0K archive/code/__torch__/caffe2 20M archive/code/__torch__/torch/fx/graph_module 20M archive/code/__torch__/torch/fx 8.0K archive/code/__torch__/torch/classes 20M archive/code/__torch__/torch 20M archive/code/__torch__ 20M archive/code 2.7M archive/constants 35M archive [qihan@devvm5585.vll0 ~]$ du resaved -h 4.0K resaved/extra 8.0K resaved/code/__torch__/caffe2/torch/fb/model_transform/splitting 8.0K resaved/code/__torch__/caffe2/torch/fb/model_transform 8.0K resaved/code/__torch__/caffe2/torch/fb 8.0K resaved/code/__torch__/caffe2/torch 8.0K resaved/code/__torch__/caffe2 1.3M resaved/code/__torch__/torch/fx/graph_module 1.3M resaved/code/__torch__/torch/fx 8.0K resaved/code/__torch__/torch/classes 1.4M resaved/code/__torch__/torch 1.4M resaved/code/__torch__ 1.4M resaved/code 2.7M resaved/constants 13M resaved [qihan@devvm5585.vll0 ~]$ ``` ## Additional test: `buck test mode/dev-tsan //caffe2/benchmarks/static_runtime:static_runtime_cpptest -- --exact 'caffe2/benchmarks/static_runtime:static_runtime_cpptest - StaticRuntime.to'` passes test jest.fbios.startup_cold_start.local.simulator f333356873 - Differential Revision: D35196883 Pull Request resolved: pytorch#74869 Approved by: https://github.com/gmagogsfm
1 parent de949a0 commit b34b192

26 files changed

+837
-236
lines changed

‎test/cpp/jit/source_range_test.cpp

+51
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
#include <gtest/gtest.h>
2+
#include <torch/csrc/jit/frontend/source_range.h>
3+
4+
using namespace ::testing;
5+
using namespace ::torch::jit;
6+
7+
TEST(SourceRangeTest, test_find) {
8+
std::vector<std::shared_ptr<std::string>> strings;
9+
strings.push_back(std::make_shared<std::string>("hello world"));
10+
strings.push_back(std::make_shared<std::string>("nihaoma"));
11+
12+
std::vector<c10::string_view> pieces{*strings[0], *strings[1]};
13+
14+
StringCordView view(pieces, strings);
15+
16+
auto x = view.find("rldni", 0);
17+
EXPECT_EQ(x, 8);
18+
}
19+
20+
TEST(SourceRangeTest, test_substr) {
21+
std::vector<std::shared_ptr<std::string>> strings;
22+
strings.push_back(std::make_shared<std::string>("hello world"));
23+
strings.push_back(std::make_shared<std::string>("nihaoma"));
24+
25+
std::vector<c10::string_view> pieces{*strings[0], *strings[1]};
26+
27+
StringCordView view(pieces, strings);
28+
29+
auto x = view.substr(4, 10).str();
30+
EXPECT_EQ(x, view.str().substr(4, 10));
31+
EXPECT_EQ(view.substr(0, view.size()).str(), view.str());
32+
}
33+
34+
TEST(SourceRangeTest, test_iter) {
35+
std::vector<std::shared_ptr<std::string>> strings;
36+
strings.push_back(std::make_shared<std::string>("hello world"));
37+
strings.push_back(std::make_shared<std::string>("nihaoma"));
38+
39+
std::vector<c10::string_view> pieces{*strings[0], *strings[1]};
40+
41+
StringCordView view(pieces, strings);
42+
43+
auto iter = view.iter_for_pos(5);
44+
EXPECT_EQ(*iter, ' ');
45+
EXPECT_EQ(iter.rest_line(), " world");
46+
EXPECT_EQ(*iter.next_iter(), 'w');
47+
EXPECT_EQ(iter.pos(), 5);
48+
49+
iter = view.iter_for_pos(13);
50+
EXPECT_EQ(iter.pos(), 13);
51+
}

‎test/cpp/jit/test_backend.cpp

+82
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,38 @@ TEST(BackendTest, TestCompiler) {
143143
AT_ASSERT(mres.toTensor().equal(ref.toTensor()));
144144
}
145145

146+
TEST(BackendTest, TestCompilerWithStringTable) {
147+
setShouldUseFormatWithStringTable(true);
148+
Module m("m");
149+
m.define(R"(
150+
def forward(self, x, h):
151+
return x + h
152+
)");
153+
154+
std::vector<IValue> inputs;
155+
inputs.emplace_back(2.0 * torch::ones({}));
156+
inputs.emplace_back(1.0 * torch::ones({}));
157+
auto ref = m.forward(inputs);
158+
159+
c10::Dict<IValue, IValue> compile_spec(StringType::get(), AnyType::get());
160+
c10::Dict<IValue, IValue> fake_dict(StringType::get(), AnyType::get());
161+
fake_dict.insert("", "");
162+
compile_spec.insert("forward", fake_dict);
163+
auto any_dict_ty = DictType::create(StringType::get(), AnyType::get());
164+
// lowered module
165+
auto lm = torch::jit::detail::codegen_backend_module(
166+
"backend_with_compiler_demo", m, compile_spec, any_dict_ty);
167+
auto res = lm.forward(inputs);
168+
AT_ASSERT(res.toTensor().equal(ref.toTensor()));
169+
170+
std::stringstream ss;
171+
lm._save_for_mobile(ss);
172+
auto mlm = _load_for_mobile(ss);
173+
auto mres = mlm.forward(inputs);
174+
setShouldUseFormatWithStringTable(false);
175+
AT_ASSERT(mres.toTensor().equal(ref.toTensor()));
176+
}
177+
146178
TEST(BackendTest, TestComposite) {
147179
c10::Dict<IValue, IValue> compile_spec(StringType::get(), AnyType::get());
148180
c10::Dict<IValue, IValue> fake_dict(StringType::get(), AnyType::get());
@@ -384,6 +416,56 @@ Traceback of TorchScript (most recent call last):
384416
ASSERT_THROWS_WITH_MESSAGE(mlm.forward(inputs), error_pattern);
385417
}
386418

419+
TEST(BackendTestDebugInfo, TestCompilerWithStringTable) {
420+
setShouldUseFormatWithStringTable(true);
421+
Module m("m");
422+
m.define(R"(
423+
def forward(self, x, h):
424+
return x + h
425+
)");
426+
427+
std::vector<IValue> inputs;
428+
inputs.emplace_back(torch::rand({2, 4}));
429+
inputs.emplace_back(torch::rand({13, 9}));
430+
431+
c10::Dict<IValue, IValue> compile_spec(StringType::get(), AnyType::get());
432+
c10::Dict<IValue, IValue> fake_dict(StringType::get(), AnyType::get());
433+
fake_dict.insert("", "");
434+
compile_spec.insert("forward", fake_dict);
435+
auto any_dict_ty = DictType::create(StringType::get(), AnyType::get());
436+
// lowered module
437+
auto lm = torch::jit::detail::codegen_backend_module(
438+
"backend_with_compiler_demo", m, compile_spec, any_dict_ty);
439+
440+
std::stringstream ss;
441+
lm._save_for_mobile(ss, ExtraFilesMap(), true);
442+
auto mlm = _load_for_mobile(ss);
443+
std::string error_pattern = R"(
444+
Module hierarchy:top(m)::<unknown>.__loweredModule__(m)::forward.aten::add
445+
Traceback of TorchScript (most recent call last):
446+
File "<string>", line 3, in <unknown>
447+
448+
def forward(self, x: Tensor, h: Tensor):
449+
return self.__loweredModule__.forward(x, h)
450+
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
451+
452+
File "<string>", line 5, in forward
453+
typed_inputs: List[Any] = [x, h, ]
454+
if self.__backend.is_available() :
455+
_0, = self.__backend.execute(self.__handles["forward"], typed_inputs)
456+
~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
457+
assert isinstance(_0, Tensor)
458+
return _0
459+
File "<string>", line 3, in <unknown>
460+
461+
def forward(self, x, h):
462+
return x + h
463+
~~~~~ <--- HERE
464+
)";
465+
setShouldUseFormatWithStringTable(false);
466+
ASSERT_THROWS_WITH_MESSAGE(mlm.forward(inputs), error_pattern);
467+
}
468+
387469
TEST(BackendTestDebugInfo, TestExceptionStackForCompilerWithModuleHierarchy) {
388470
Module a("A");
389471
a.define(R"(

‎test/cpp/jit/test_utils.h

+12-10
Original file line numberDiff line numberDiff line change
@@ -32,16 +32,18 @@ static inline void trim(std::string& s) {
3232
}
3333
} // namespace
3434

35-
#define ASSERT_THROWS_WITH_MESSAGE(statement, substring) \
36-
try { \
37-
(void)statement; \
38-
FAIL(); \
39-
} catch (const std::exception& e) { \
40-
std::string substring_s(substring); \
41-
trim(substring_s); \
42-
auto exception_string = std::string(e.what()); \
43-
trim(exception_string); \
44-
ASSERT_NE(exception_string.find(substring_s), std::string::npos); \
35+
#define ASSERT_THROWS_WITH_MESSAGE(statement, substring) \
36+
try { \
37+
(void)statement; \
38+
FAIL(); \
39+
} catch (const std::exception& e) { \
40+
std::string substring_s(substring); \
41+
trim(substring_s); \
42+
auto exception_string = std::string(e.what()); \
43+
trim(exception_string); \
44+
ASSERT_NE(exception_string.find(substring_s), std::string::npos) \
45+
<< " Error was: \n" \
46+
<< exception_string; \
4547
}
4648

4749
namespace torch {

‎torch/csrc/jit/frontend/function_schema_parser.cpp

+8-2
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
#include <ATen/core/Reduction.h>
44
#include <ATen/core/type_factory.h>
5+
#include <c10/util/Optional.h>
56
#include <c10/util/string_utils.h>
67
#include <torch/csrc/jit/frontend/lexer.h>
78
#include <torch/csrc/jit/frontend/parse_string_literal.h>
@@ -27,8 +28,13 @@ namespace jit {
2728

2829
namespace {
2930
struct SchemaParser {
30-
SchemaParser(const std::string& str)
31-
: L(std::make_shared<SourceView>(c10::string_view(str))),
31+
explicit SchemaParser(const std::string& str)
32+
: L(std::make_shared<Source>(
33+
c10::string_view(str),
34+
c10::nullopt,
35+
0,
36+
nullptr,
37+
Source::DONT_COPY)),
3238
type_parser(L, /*parse_complete_tensor_types*/ false) {}
3339

3440
either<OperatorName, FunctionSchema> parseDeclaration() {

0 commit comments

Comments
 (0)