已合并
临时跳过UT用例test_output_chunks_memory_usage #39231
dairenjie创建于 6月24日
临时跳过UT用例test_output_chunks_memory_usage #39231
已合并
共 1 个文件变更+2-0
| @@ -12,6 +12,7 @@ | |||
| 12 | # See the License for the specific language governing permissions and | 12 | # See the License for the specific language governing permissions and |
| 13 | # limitations under the License. | 13 | # limitations under the License. |
| 14 | import os | 14 | import os |
| 15 | +import unittest | ||
| 15 | 16 | ||
| 16 | from model_registry import ExampleCode, ModelWithKwargs, MultiMLP | 17 | from model_registry import ExampleCode, ModelWithKwargs, MultiMLP |
| 17 | 18 | ||
| @@ -270,6 +271,7 @@ class StageTest(MultiProcContinuousTest): | |||
| 270 | ref_out = full_mod(x) | 271 | ref_out = full_mod(x) |
| 271 | torch.testing.assert_close(out, ref_out) | 272 | torch.testing.assert_close(out, ref_out) |
| 272 | 273 | ||
| 274 | + | ||
| 273 | def test_output_chunks_memory_usage(self): | 275 | def test_output_chunks_memory_usage(self): |
| 274 | """Test that output_chunks doesn't store memory for non-first stages.""" | 276 | """Test that output_chunks doesn't store memory for non-first stages.""" |
| 275 | full_mod = MultiMLP(d_hid, n_layers=self.world_size) | 277 | full_mod = MultiMLP(d_hid, n_layers=self.world_size) |