diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index 9a214964ce1..9894a305254 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -716,7 +716,11 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then # executor drives, plus the vision encoder and forward_from_embeddings. OFFGRAPH_FLAG="" if [ "$MODE" = "dflash-image-batching" ]; then - OFFGRAPH_FLAG="--use-offgraph-kv-cache" + # Every static forward_{N} and draft_forward_{N} is its own AOTI + # compile and autotune, which the 11-vCPU runner cannot fit in the + # job timeout at the default sizes; these still cover static and + # dynamic target steps and DFlash drafting. + OFFGRAPH_FLAG="--use-offgraph-kv-cache --forward-sizes 1,4 --max-draft-owners 1" fi EXPORT_START_SECONDS=$SECONDS TMPDIR="$INDUCTOR_TMPDIR" \ diff --git a/.ci/scripts/tests/test_muse_glimmer_batching_ci.py b/.ci/scripts/tests/test_muse_glimmer_batching_ci.py index 5c8fcead860..b04eb9eab23 100644 --- a/.ci/scripts/tests/test_muse_glimmer_batching_ci.py +++ b/.ci/scripts/tests/test_muse_glimmer_batching_ci.py @@ -129,6 +129,26 @@ def test_smoke_cmake_is_opt_in_and_reuses_runtime(self): self.assertNotIn("BENCHMARK", smoke) self.assertNotIn("batching_bench", smoke) + def test_full_e2e_runs_only_on_main_pushes(self): + workflow = (ROOT / ".github/workflows/cuda.yml").read_text() + e2e_if = ( + workflow.split(" test-muse-glimmer-cuda-e2e:", 1)[1] + .split(" if: |\n", 1)[1] + .split(" uses:", 1)[0] + .strip() + ) + self.assertTrue( + e2e_if.startswith( + "github.event_name == 'push' && github.ref == 'refs/heads/main' && (" + ) + ) + self.assertTrue(e2e_if.endswith(")")) + self.assertEqual(e2e_if.count("("), e2e_if.count(")")) + unit_if = workflow.split(" unittest-cuda-muse-glimmer-batching:", 1)[1].split( + " uses:", 1 + )[0] + self.assertNotIn("github.event_name", unit_if) + def test_owned_shell_and_unit_workflow_syntax(self): for script in ("test_model_e2e.sh", "build_muse_glimmer_batching_smoke.sh"): result = subprocess.run(["bash", "-n", str(SCRIPTS / script)]) diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 56bf2f71262..cfc9006f557 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -1197,7 +1197,11 @@ jobs: test-muse-glimmer-cuda-e2e: name: test-muse-glimmer-cuda-e2e-${{ matrix.variant }}-${{ matrix.mode }} needs: [changed-files, run-decision] + # Hours of A100 time per run, so only after merge to main (not on pull + # requests, ciflow tags or manual dispatch); pull requests get the + # unittest-cuda-* jobs above. if: | + ( contains(needs.changed-files.outputs.changed-files, '.ci/scripts/build_muse_glimmer_batching_smoke.sh') || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/muse_glimmer_batching_smoke.cpp') || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/muse_glimmer_batching_workload.h') || @@ -1212,7 +1216,7 @@ jobs: contains(needs.changed-files.outputs.changed-files, '.ci/scripts/export_model_artifact.sh') || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e.sh') || contains(needs.changed-files.outputs.changed-files, 'examples/models/muse-glimmer') || - needs.run-decision.outputs.is-full-run == 'true' + needs.run-decision.outputs.is-full-run == 'true') uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 5ab15c33016..68e3777a06f 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -3,6 +3,18 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") +fbcode_target( + _kind = runtime.python_library, + name = "cpu_count", + srcs = [ + "utils/__init__.py", + "utils/cpu_count.py", + ], + visibility = [ + "//executorch/backends/cuda/...", + ], +) + fbcode_target( _kind = runtime.python_library, name = "coalesced_int4_tensor", @@ -188,6 +200,7 @@ fbcode_target( ], deps = [ ":autotune", + ":cpu_count", "//caffe2:torch", ], ) diff --git a/backends/cuda/autotune/cuda_graph_timing.py b/backends/cuda/autotune/cuda_graph_timing.py index 66a969ea8f2..bba332a48ae 100644 --- a/backends/cuda/autotune/cuda_graph_timing.py +++ b/backends/cuda/autotune/cuda_graph_timing.py @@ -31,8 +31,11 @@ captured). To keep captures succeeding: -- each capture's graph and its memory pool are released right after timing, - and the allocator's cache is emptied once it holds a lot of unused memory; +- each capture's graph is released and the allocator's cache emptied right + after timing. The cache is not checked first: the allocator stats keep an + entry for every private pool ever captured, so reading them costs more with + each capture, while emptying a cache that holds little more than the + released pool is cheap; - on CUDA OOM the cache is emptied and the capture retried; if that fails too, the tensors given as ``offload`` (the model's weights moved to the GPU for the compile) are parked on the CPU while the candidate is timed and restored @@ -59,7 +62,6 @@ _TARGET_GRAPH_US = 2000.0 _MIN_CALLS = 4 _MAX_CALLS = 32 -_EMPTY_CACHE_THRESHOLD_BYTES = 1 << 30 @dataclass @@ -116,14 +118,6 @@ def _default_cuda_bench( return bench.benchmark_gpu(fn, warmup=warmup, rep=rep, **kwargs) -def _release_cached_memory_if_large() -> None: - if ( - torch.cuda.memory_reserved() - torch.cuda.memory_allocated() - > _EMPTY_CACHE_THRESHOLD_BYTES - ): - torch.cuda.empty_cache() - - _l2_flush_buffers: dict[int, torch.Tensor] = {} @@ -180,7 +174,7 @@ def time_with_cuda_graph_us(fn: Callable[[], Any]) -> float: with contextlib.suppress(RuntimeError): graph.reset() del graph - _release_cached_memory_if_large() + torch.cuda.empty_cache() class ParkedTensorsNotRestored(RuntimeError): diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 64f0866a642..3e9097838a8 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -1166,7 +1166,7 @@ def materialize_weights_blob( blob_path = os.path.splitext(so_path)[0] + "_weights.blob" capture = CudaWeightCollector.current_capture() capture.artifact = capture.collector.materialize( - weights[0], os.path.dirname(blob_path), capture.device_type_for_weight + weights[0], os.path.dirname(blob_path), capture ) # Keep AotiBackend's existing path contract intact. The compatibility @@ -1545,14 +1545,30 @@ def move_program_to_device( check on the cache update. They are re-synthesized as zeros for the lifted graph and the serialized blob. This activates automatically with low-memory mode. Other (non-low-memory) exports use the stock pass. + + Low-memory methods share their CPU weights, so the CPU source of each + moved tensor is recorded on the active weight capture. ``materialize`` + then reuses a storage an earlier method already wrote from the same + source view. The moved storages must keep their source's bytes until + ``materialize`` runs. """ from torch.export.passes import move_to_device_pass if not cls._is_low_memory_mode(compile_specs): moved = move_to_device_pass(edge_program, device) - else: - moved = _move_to_device_resize_kv(edge_program, device) + _track_moved_tensors(moved) + return moved + # Only parameters are recorded: low-memory compilation can hand AOTI CPU + # clones of buffers (_compile_time_cpu_clones). + sources = { + name: edge_program.state_dict[name] + for name in edge_program.graph_signature.parameters + } + moved = _move_to_device_resize_kv(edge_program, device) _track_moved_tensors(moved) + CudaWeightCollector.current_capture().record_moved_sources( + sources, dict(moved.state_dict) + ) return moved @classmethod diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index f826a34d718..c36676a33da 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -21,6 +21,7 @@ from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.backend_details import PreprocessResult from executorch.exir.tensor import scalar_type_enum +from torch.multiprocessing.reductions import StorageWeakRef CUDA_AOTI_METADATA_MAGIC = b"ETCUDAFQN0" @@ -46,7 +47,9 @@ class CudaWeightArtifact: entries: List[CudaWeightEntry] storages: Dict[str, FileBackedData] # Storage keys an earlier method of this collector already stored with the - # same bytes; they are not written again. + # same bytes; they are not written again. A key reused because it was moved + # from the same CPU source view is not re-read: that relies on the source + # not changing between methods in ways its version counter misses. stored: Set[str] = field(default_factory=set) @@ -68,11 +71,71 @@ class CudaAotiMetadata: entries: List[CudaWeightEntry] +def _source_identity(source: torch.Tensor) -> Tuple[Any, ...]: + # Bumped by in-place ops autograd tracks; writes through .data, .numpy() or + # the storage do not bump it. + return ( + source.untyped_storage()._cdata, + source._version, + source.dtype, + source.storage_offset(), + tuple(source.shape), + tuple(source.stride()), + ) + + +@dataclass(frozen=True) +class _MovedSource: + source: StorageWeakRef + identity: Tuple[Any, ...] + # The moved tensor's version when it was recorded; a write to it before + # materialize makes it be read and checked again. + moved_version: int + + @dataclass class _CudaWeightCapture: collector: "CudaWeightCollector" device_type_for_weight: Callable[[torch.Tensor], int] artifact: Optional[CudaWeightArtifact] = None + # This method's moved storage -> the CPU source it was copied from. The + # StorageWeakRefs (which hash and compare by address) keep both storage + # objects, not their data, allocated, so no other storage can take either + # address while this capture lives. + moved_sources: Dict[StorageWeakRef, "_MovedSource"] = field(default_factory=dict) + + def record_moved_sources( + self, sources: Dict[str, torch.Tensor], moved: Dict[str, torch.Tensor] + ) -> None: + """Record the CPU tensor each tensor of this method was moved from. + + Must run inside ``capture()`` before ``materialize``. A matched storage + is reused without being read, so nothing may write the moved tensors + between this call and ``materialize`` except through in-place ops + their version counter records. + """ + self.moved_sources = {} + for name, tensor in moved.items(): + source = sources.get(name) + if ( + type(source) in (torch.Tensor, torch.nn.Parameter) + and type(tensor) in (torch.Tensor, torch.nn.Parameter) + and tensor.is_cuda + and source.device.type == "cpu" + # Weights loaded under inference_mode can change without a + # version bump. + and not source.is_inference() + and source.untyped_storage().nbytes() > 0 + # Low-memory KV placeholders are emptied on the device. + and tensor.untyped_storage().nbytes() > 0 + ): + self.moved_sources[StorageWeakRef(tensor.untyped_storage())] = ( + _MovedSource( + source=StorageWeakRef(source.untyped_storage()), + identity=_source_identity(source), + moved_version=tensor._version, + ) + ) def trim_host_memory() -> None: @@ -384,13 +447,20 @@ class CudaWeightCollector: def __init__(self) -> None: self._store = NamedDataStore() self._results: List[PreprocessResult] = [] - # The CUDA storages already written to the store, by storage key. Methods - # share their weights, so a later method's identical storage is compared - # with the first one on the device and not written again. The weights - # must not change until finish() releases them. + # The CUDA storages already written to the store, by storage key. A later + # method's identical storage is compared with the first one on the + # device and not written again while the first one is still allocated. self._stored_storages: Dict[str, torch.Tensor] = {} # The key each storage was written under, which ends in its SHA-256. self._content_keys: Dict[str, str] = {} + # The CPU source view each storage key was written from. When methods + # share their CPU weights (low_memory_mode), each moves the same source + # to the GPU and frees the copy afterwards, so a later method's copy of + # the same source view is recognized by its source and not written + # again. The StorageWeakRef keeps the source's storage object allocated, + # so its address, part of the identity, is not reused; once its data is + # freed, a later live source takes the key over. + self._written_sources: Dict[str, _MovedSource] = {} @contextlib.contextmanager def capture( @@ -417,9 +487,14 @@ def materialize( self, weights: Any, directory: str, - device_type_for_weight: Callable[[torch.Tensor], int], + capture: _CudaWeightCapture, ) -> CudaWeightArtifact: - """Turn AOTI ``Weights`` into one independently named blob per FQN.""" + """Turn AOTI ``Weights`` into one independently named blob per FQN. + + ``capture`` holds the sources ``move_program_to_device`` recorded for + this method; storages copied from an already written source view are + reused without being read. + """ trim_host_memory() entries: List[CudaWeightEntry] = [] storages: Dict[str, FileBackedData] = {} @@ -440,7 +515,7 @@ def materialize( and tuple(tensor.stride()) == tuple(getattr(properties, "stride", ())) ) del storage - device_type = device_type_for_weight(tensor) + device_type = capture.device_type_for_weight(tensor) expected_storage_nbytes = int( getattr(properties, "storage_size", None) or 0 ) @@ -460,9 +535,24 @@ def materialize( raise RuntimeError(f"Duplicate CUDA FQN weight key for {fqn!r}") # Library-local constants are keyed per shared library, so # only FQN weights can match an earlier method's storage. - shareable = tensor.is_cuda and not _is_aoti_library_local_fqn(fqn) + library_local = _is_aoti_library_local_fqn(fqn) + shareable = tensor.is_cuda and not library_local first = self._stored_storages.get(storage_key) + moved = ( + capture.moved_sources.get(StorageWeakRef(tensor.untyped_storage())) + if shareable + else None + ) + if moved is not None and tensor._version != moved.moved_version: + moved = None + written = self._written_sources.get(storage_key) if ( + moved is not None + and written is not None + and written.identity == moved.identity + ): + stored.add(storage_key) + elif ( shareable and first is not None and _same_cuda_storage(first, tensor) @@ -472,6 +562,10 @@ def materialize( storages[storage_key] = _materialize_storage(tensor, directory) if shareable: self._stored_storages.setdefault(storage_key, tensor) + if moved is not None and ( + written is None or written.source.expired() + ): + self._written_sources[storage_key] = moved sizes = tuple( int(size) for size in getattr(properties, "shape", tensor.shape) @@ -639,3 +733,4 @@ def finish(self) -> None: result.data_store_output = shared_output self._stored_storages.clear() self._content_keys.clear() + self._written_sources.clear() diff --git a/backends/cuda/tests/autotune_test_utils.py b/backends/cuda/tests/autotune_test_utils.py index 3e49bff3f94..43b1e677bd2 100644 --- a/backends/cuda/tests/autotune_test_utils.py +++ b/backends/cuda/tests/autotune_test_utils.py @@ -17,6 +17,7 @@ import torch import triton from executorch.backends.cuda.autotune import cuda_graph_timing as cgt +from executorch.backends.cuda.utils.cpu_count import available_cpu_count from torch._inductor.runtime.triton_heuristics import CachingAutotuner from torch._inductor.select_algorithm import AlgorithmSelectorCache @@ -58,13 +59,13 @@ def _spin(stop_at: float) -> None: class SaturatedCpu: - """One busy-loop process per core; can be paused to measure a reference.""" + """One busy-loop process per usable CPU; can be paused to measure a reference.""" def __init__(self) -> None: stop_at = time.time() + 3600 self.workers = [ multiprocessing.Process(target=_spin, args=(stop_at,), daemon=True) - for _ in range(os.cpu_count() or 1) + for _ in range(available_cpu_count()) ] def __enter__(self) -> "SaturatedCpu": diff --git a/backends/cuda/tests/targets.bzl b/backends/cuda/tests/targets.bzl index f403dfb60ad..09b3700cbf5 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -7,6 +7,16 @@ def define_common_targets(is_fbcode = False): if not is_fbcode: return + python_unittest( + name = "test_cpu_count", + srcs = [ + "test_cpu_count.py", + ], + deps = [ + "//executorch/backends/cuda:cpu_count", + ], + ) + python_unittest_remote_gpu( name = "test_cuda_export", srcs = [ @@ -42,6 +52,7 @@ def define_common_targets(is_fbcode = False): deps = [ "//caffe2:torch", "//executorch/backends/cuda:autotune", + "//executorch/backends/cuda:cpu_count", ], ) @@ -383,6 +394,7 @@ def define_common_targets(is_fbcode = False): deps = [ "//caffe2:torch", "//executorch/backends/cuda:cuda_backend", + "//executorch/backends/cuda:cuda_partitioner", "//executorch/exir:lib", ], keep_gpu_sections = True, diff --git a/backends/cuda/tests/test_cpu_count.py b/backends/cuda/tests/test_cpu_count.py new file mode 100644 index 00000000000..50f040abf43 --- /dev/null +++ b/backends/cuda/tests/test_cpu_count.py @@ -0,0 +1,147 @@ +# Copyright (c) Meta Platforms, Inc. and 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. + +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +from executorch.backends.cuda.utils import cpu_count + + +class AvailableCpuCountTest(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.mount = Path(self.temporary.name) + self.proc_cgroup = "0::/job\n" + self.proc_mountinfo = f"1 0 0:1 / {self.mount} rw - cgroup2 cgroup rw\n" + original_read = Path.read_text + + def read_text(path, *args, **kwargs): + if path == Path("/proc/self/cgroup"): + return self.proc_cgroup + if path == Path("/proc/self/mountinfo"): + return self.proc_mountinfo + return original_read(path, *args, **kwargs) + + for patch in ( + mock.patch.object(cpu_count.sys, "platform", "linux"), + mock.patch.object(cpu_count.os, "cpu_count", return_value=44), + mock.patch.object( + cpu_count.os, + "sched_getaffinity", + return_value=set(range(44)), + create=True, + ), + mock.patch.object(Path, "read_text", read_text), + ): + patch.start() + self.addCleanup(patch.stop) + + def quota(self, group, quota, period=100000, unified=True): + directory = self.mount / group + directory.mkdir(parents=True, exist_ok=True) + if unified: + (directory / "cpu.max").write_text(f"{quota} {period}\n") + else: + (directory / "cpu.cfs_quota_us").write_text(str(quota)) + (directory / "cpu.cfs_period_us").write_text(str(period)) + + def test_unlimited_host_preserves_cpu_count(self): + self.quota("job", "max") + self.assertEqual(cpu_count.available_cpu_count(), 44) + + def test_affinity_restricts_available_cpus(self): + self.quota("job", "max") + with mock.patch.object( + cpu_count.os, "sched_getaffinity", return_value=set(range(7)) + ): + self.assertEqual(cpu_count.available_cpu_count(), 7) + + def test_cgroup_v2_quota_without_affinity_limit(self): + self.quota("job", 1100000) + self.assertEqual(cpu_count.available_cpu_count(), 11) + + def test_tightest_visible_parent_quota(self): + self.proc_cgroup = "0::/parent/job\n" + self.quota("parent/job", "max") + self.quota("parent", 1100000) + self.quota("", 800000) + self.assertEqual(cpu_count.available_cpu_count(), 8) + + def test_mount_root_maps_host_cgroup_path(self): + self.proc_cgroup = "0::/containers/container/job\n" + self.proc_mountinfo = ( + f"1 0 0:1 /containers/container {self.mount} rw - cgroup2 cgroup rw\n" + ) + self.quota("job", 1100000) + self.assertEqual(cpu_count.available_cpu_count(), 11) + + def test_cgroup_namespace_root_maps_mountpoint(self): + self.proc_cgroup = "0::/\n" + self.proc_mountinfo = ( + f"1 0 0:1 /containers/container {self.mount} rw - cgroup2 cgroup rw\n" + ) + self.quota("", 1100000) + self.assertEqual(cpu_count.available_cpu_count(), 11) + + def test_cgroup_v1_combined_controllers(self): + self.proc_cgroup = "4:cpu,cpuacct:/job\n2:memory:/other\n" + self.proc_mountinfo = ( + f"1 0 0:1 / {self.mount} rw - cgroup cgroup rw,cpu,cpuacct\n" + ) + self.quota("job", 1100000, unified=False) + self.assertEqual(cpu_count.available_cpu_count(), 11) + + def test_cgroup_v1_unlimited_parent_can_restrict(self): + self.proc_cgroup = "4:cpu:/parent/job\n" + self.proc_mountinfo = f"1 0 0:1 / {self.mount} rw - cgroup cgroup rw,cpu\n" + self.quota("parent/job", -1, unified=False) + self.quota("parent", 1100000, unified=False) + self.assertEqual(cpu_count.available_cpu_count(), 11) + + def test_fractional_quota_rounds_up_to_a_worker(self): + for quota, expected in ((50000, 1), (150000, 2)): + with self.subTest(quota=quota): + self.quota("job", quota) + self.assertEqual(cpu_count.available_cpu_count(), expected) + + def test_affinity_tighter_than_quota(self): + self.quota("job", 1100000) + with mock.patch.object(cpu_count.os, "sched_getaffinity", return_value={0, 1}): + self.assertEqual(cpu_count.available_cpu_count(), 2) + + def test_missing_proc_files_falls_back(self): + with mock.patch.object(Path, "read_text", side_effect=PermissionError): + self.assertEqual(cpu_count.available_cpu_count(), 44) + + def test_invalid_or_missing_quota_does_not_hide_parent(self): + self.proc_cgroup = "0::/parent/job\n" + self.quota("parent", 1100000) + for value in (None, "invalid", "0 100000", "1100000 0"): + with self.subTest(value=value): + if value is not None: + self.quota("parent/job", 1) + (self.mount / "parent/job/cpu.max").write_text(value) + self.assertEqual(cpu_count.available_cpu_count(), 11) + + def test_non_linux_without_affinity(self): + with ( + mock.patch.object(cpu_count.sys, "platform", "win32"), + mock.patch.object( + cpu_count.os, "sched_getaffinity", side_effect=AttributeError + ), + ): + self.assertEqual(cpu_count.available_cpu_count(), 44) + + def test_missing_cpu_count_still_returns_one(self): + with mock.patch.object(cpu_count.os, "cpu_count", return_value=None): + self.assertEqual(cpu_count.available_cpu_count(), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/cuda/tests/test_cuda_graph_timing.py b/backends/cuda/tests/test_cuda_graph_timing.py index 5cf2f4b5f75..f0edc318307 100644 --- a/backends/cuda/tests/test_cuda_graph_timing.py +++ b/backends/cuda/tests/test_cuda_graph_timing.py @@ -153,6 +153,22 @@ def test_compile_spec_turns_it_off(self) -> None: self.assertIs(benchmarking._BENCHMARK_DISPATCH.get("cuda"), before) self.assertIs(benchmarking._BENCHMARK_DISPATCH.get("cuda"), before) + def test_captures_do_not_read_the_allocator_stats(self) -> None: + x = torch.zeros(1 << 16, device="cuda") + # The torch.cuda and torch.accelerator stats APIs all end in these. + with mock.patch.object( + torch._C, "_cuda_memoryStats", wraps=torch._C._cuda_memoryStats + ) as cuda_stats, mock.patch.object( + torch._C, + "_accelerator_getDeviceStats", + wraps=torch._C._accelerator_getDeviceStats, + ) as accelerator_stats, cgt.cuda_graph_timing() as stats: + for _ in range(8): + benchmarking.benchmarker.benchmark(lambda: x.add_(1), device="cuda") + self.assertEqual(stats.captured, 8) + self.assertEqual(cuda_stats.call_count, 0) + self.assertEqual(accelerator_stats.call_count, 0) + # --------------------------------------------------------------------------- # End to end: export while the CPU is saturated, and check every pick. diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index 159fdeeaa1f..337e1f53d72 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -50,9 +50,9 @@ class TestCudaLowMemoryExport(unittest.TestCase): @staticmethod def _materialize(weights, directory, device_type=AOTI_DEVICE_TYPE_CPU): - return CudaWeightCollector().materialize( - weights, directory, lambda _: device_type - ) + collector = CudaWeightCollector() + with collector.capture(lambda _: device_type) as capture: + return collector.materialize(weights, directory, capture) @staticmethod def _parent_result(so_key: str) -> PreprocessResult: @@ -222,9 +222,9 @@ def test_low_memory_cpu_clones_keep_cuda_device_type(self, _) -> None: weights = Weights({"weight": (tensor, TensorProperties(tensor))}) with tempfile.TemporaryDirectory() as directory: - artifact = CudaWeightCollector().materialize( - weights, directory, _aoti_device_type_for_weight - ) + collector = CudaWeightCollector() + with collector.capture(_aoti_device_type_for_weight) as capture: + artifact = collector.materialize(weights, directory, capture) self.assertEqual(artifact.entries[0].device_type, AOTI_DEVICE_TYPE_CUDA) for storage in artifact.storages.values(): storage.close() diff --git a/backends/cuda/tests/test_cuda_weight_reuse.py b/backends/cuda/tests/test_cuda_weight_reuse.py index 4413ecfd0f5..7e1fadc0af8 100644 --- a/backends/cuda/tests/test_cuda_weight_reuse.py +++ b/backends/cuda/tests/test_cuda_weight_reuse.py @@ -6,6 +6,8 @@ """CudaWeightCollector writes a weight storage the methods share only once.""" +import contextlib +import gc import tempfile import unittest from unittest import mock @@ -40,6 +42,20 @@ def _as_bytes(buffer) -> bytes: return buffer if isinstance(buffer, bytes) else buffer.to_bytes() +@contextlib.contextmanager +def _counting_writes(): + """Yields the sizes of the storages the collector writes.""" + writes = [] + materialize_storage = cuda_weight_collector._materialize_storage + + def counting(tensor, directory): + writes.append(tensor.untyped_storage().nbytes()) + return materialize_storage(tensor, directory) + + with mock.patch.object(cuda_weight_collector, "_materialize_storage", counting): + yield writes + + class CudaWeightReuseTest(unittest.TestCase): def setUp(self) -> None: if not torch.cuda.is_available(): @@ -49,9 +65,8 @@ def setUp(self) -> None: def _method(self, collector, so_key, weights): """One method's preprocess: materialize, then add to the shared store.""" - artifact = collector.materialize( - weights, self.directory.name, lambda _: AOTI_DEVICE_TYPE_CUDA - ) + with collector.capture(lambda _: AOTI_DEVICE_TYPE_CUDA) as capture: + artifact = collector.materialize(weights, self.directory.name, capture) result = _parent_result(so_key) collector.add_preprocess_result(result, artifact, "cuda") return artifact, result @@ -60,14 +75,7 @@ def test_a_storage_shared_by_methods_is_written_once(self) -> None: weight = torch.randn(1024, 33, device="cuda") bias = torch.arange(7, dtype=torch.int8, device="cuda") collector = CudaWeightCollector() - writes = [] - materialize_storage = cuda_weight_collector._materialize_storage - - def counting(tensor, directory): - writes.append(tensor.untyped_storage().nbytes()) - return materialize_storage(tensor, directory) - - with mock.patch.object(cuda_weight_collector, "_materialize_storage", counting): + with _counting_writes() as writes: first, first_result = self._method( collector, "so-0", _weights(w=weight, b=bias) ) @@ -169,3 +177,185 @@ def test_same_cuda_storage_rejects_other_sizes_and_empty_storages(self) -> None: self.assertFalse(same(base, base[:8].clone())) empty = torch.empty(0, dtype=torch.uint8, device="cuda") self.assertTrue(same(empty, empty.clone())) + + def _moved_method(self, collector, so_key, sources, replace=None, update=None): + """A method whose weights are moved from ``sources`` and freed after it. + + ``replace`` drops the moved copies before AOTI sees them (as a pass + replacing them would) and returns the tensors to materialize instead. + ``update`` writes to the moved copies before AOTI sees them. + """ + moved = {name: tensor.to("cuda") for name, tensor in sources.items()} + with collector.capture(lambda _: AOTI_DEVICE_TYPE_CUDA) as capture: + capture.record_moved_sources(sources, moved) + if update is not None: + with torch.no_grad(): + for tensor in moved.values(): + update(tensor) + if replace is not None: + shapes = {name: tensor.shape for name, tensor in moved.items()} + # Only the capture's records may keep the storages alive. + moved.clear() + gc.collect() + moved = replace(shapes) + artifact = collector.materialize( + _weights(**moved), self.directory.name, capture + ) + result = _parent_result(so_key) + collector.add_preprocess_result(result, artifact, "cuda") + for tensor in moved.values(): + tensor.untyped_storage().resize_(0) + return artifact, result + + def test_freed_copies_of_one_source_are_written_once(self) -> None: + sources = { + "w": torch.randn(1024, 33), + "b": torch.arange(7, dtype=torch.int8), + } + collector = CudaWeightCollector() + with _counting_writes() as writes: + first, first_result = self._moved_method(collector, "so-0", sources) + second, second_result = self._moved_method(collector, "so-1", sources) + third, _ = self._moved_method(collector, "so-2", sources) + collector.finish() + self.assertEqual(len(writes), 2) + self.assertEqual(second.storages, {}) + self.assertEqual(third.storages, {}) + self.assertEqual(second.stored, {e.storage_key for e in second.entries}) + self.assertEqual( + decode_cuda_aoti_metadata(second_result.processed_bytes).entries, + decode_cuda_aoti_metadata(first_result.processed_bytes).entries, + ) + stored = { + key: _as_bytes(collector._store.buffers[index]) + for key, index in collector._store.key_to_buffer_idx.items() + } + for entry in decode_cuda_aoti_metadata(first_result.processed_bytes).entries: + self.assertEqual( + stored[entry.storage_key], + sources[entry.fqn].contiguous().view(torch.uint8).numpy().tobytes(), + ) + + def test_another_source_with_the_same_key_is_written_and_checked(self) -> None: + collector = CudaWeightCollector() + self._moved_method(collector, "so-0", {"w": torch.randn(256, 8)}) + with self.assertRaisesRegex(ValueError, "with different data"): + self._moved_method(collector, "so-1", {"w": torch.randn(256, 8)}) + + def test_a_source_view_with_other_metadata_is_written_again(self) -> None: + base = torch.randn(64, 16) + for first, second in ( + (base, base[:32]), # shape + (base[:32], base[32:]), # storage offset + (base[:, :8], base[:, ::2]), # strides + ): + with self.subTest(second=(second.shape, second.storage_offset())): + collector = CudaWeightCollector() + with _counting_writes() as writes: + self._moved_method(collector, "so-0", {"w": first}) + with self.assertRaisesRegex(ValueError, "with different data"): + self._moved_method(collector, "so-1", {"w": second}) + self.assertEqual(len(writes), 2) + + def test_a_source_updated_in_place_is_written_and_checked(self) -> None: + sources = {"w": torch.nn.Parameter(torch.randn(256, 8))} + collector = CudaWeightCollector() + self._moved_method(collector, "so-0", sources) + with torch.no_grad(): + sources["w"].add_(1) + with self.assertRaisesRegex(ValueError, "with different data"): + self._moved_method(collector, "so-1", sources) + + def test_a_moved_copy_written_before_materialize_is_checked(self) -> None: + sources = {"w": torch.randn(256, 8)} + collector = CudaWeightCollector() + self._moved_method(collector, "so-0", sources) + with self.assertRaisesRegex(ValueError, "with different data"): + self._moved_method( + collector, "so-1", sources, update=lambda tensor: tensor.add_(1) + ) + + def test_a_storage_freed_before_materialize_is_not_matched(self) -> None: + sources = {"w": torch.randn(256, 8)} + collector = CudaWeightCollector() + self._moved_method(collector, "so-0", sources) + # Same source, but the moved copy is freed before AOTI sees it and a + # storage with other bytes is allocated, possibly at its old address: + # it must be read and checked. + with self.assertRaisesRegex(ValueError, "with different data"): + self._moved_method( + collector, + "so-1", + sources, + replace=lambda shapes: { + name: torch.randn(shape, device="cuda") + for name, shape in shapes.items() + }, + ) + + def test_a_freed_written_source_lets_a_live_source_take_over(self) -> None: + collector = CudaWeightCollector() + first = torch.randn(256, 8) + with _counting_writes() as writes: + self._moved_method(collector, "so-0", {"w": first}) + shared = first.clone() + del first + self._moved_method(collector, "so-1", {"w": shared}) + self._moved_method(collector, "so-2", {"w": shared}) + self.assertEqual(len(writes), 2) + + +class _TwoMethodLinear(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(64, 48) + # A persistent buffer keeps the per-method handling. + self.register_buffer("scale", (torch.rand(48) + 0.5).to(torch.bfloat16)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) * self.scale + + +class CudaWeightReuseExportTest(unittest.TestCase): + """Lowers two methods through CudaBackend, as a real export does.""" + + def setUp(self) -> None: + if not torch.cuda.is_available(): + self.skipTest("CUDA required") + + def _lower(self, low_memory: bool): + from executorch.backends.cuda.cuda_backend import CudaBackend + from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower + from executorch.exir.backend.compile_spec_schema import CompileSpec + + module = _TwoMethodLinear().eval() + programs = { + name: torch.export.export(module, (torch.randn(rows, 64),), strict=True) + for name, rows in (("forward_1", 1), ("forward_4", 4)) + } + + def partitioner(name): + specs = [CudaBackend.generate_method_name_compile_spec(name)] + if low_memory: + specs.append(CompileSpec("low_memory_mode", b"ON")) + return [CudaPartitioner(specs)] + + with _counting_writes() as writes: + to_edge_transform_and_lower( + programs, + partitioner={name: partitioner(name) for name in programs}, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + return writes + + # The linear weight, its bias and the bf16 buffer. + _SIZES = (64 * 48 * 4, 48 * 4, 48 * 2) + + def test_low_memory_methods_write_shared_weights_once(self) -> None: + writes = self._lower(low_memory=True) + self.assertEqual([writes.count(size) for size in self._SIZES], [1, 1, 2]) + + def test_default_mode_writes_each_methods_own_copy(self) -> None: + writes = self._lower(low_memory=False) + self.assertEqual([writes.count(size) for size in self._SIZES[:2]], [2, 2]) diff --git a/backends/cuda/triton/kernels/quantized_gemm_family.py b/backends/cuda/triton/kernels/quantized_gemm_family.py index aa58b17699c..3a30a031404 100644 --- a/backends/cuda/triton/kernels/quantized_gemm_family.py +++ b/backends/cuda/triton/kernels/quantized_gemm_family.py @@ -42,6 +42,7 @@ import torch import triton from executorch.backends.cuda.triton.kernels.quantized_gemm_utils import splitk_reduce +from executorch.backends.cuda.utils.cpu_count import available_cpu_count from torch.library import triton_op, wrap_triton try: @@ -180,7 +181,7 @@ def _has_untuned_configs(kernel, args: Sequence, kwargs: dict) -> bool: def _compile_workers() -> int: - return min(_MAX_COMPILE_THREADS, os.cpu_count() or 1) + return min(_MAX_COMPILE_THREADS, available_cpu_count()) _PROCESS_POOL_LOCK = threading.Lock() diff --git a/backends/cuda/utils/__init__.py b/backends/cuda/utils/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/cuda/utils/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and 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. diff --git a/backends/cuda/utils/cpu_count.py b/backends/cuda/utils/cpu_count.py new file mode 100644 index 00000000000..b03f6599f0e --- /dev/null +++ b/backends/cuda/utils/cpu_count.py @@ -0,0 +1,108 @@ +# Copyright (c) Meta Platforms, Inc. and 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. + +import os +import re +import sys +from pathlib import Path, PurePosixPath + + +def _cgroup_memberships(memberships: list[str]) -> dict[str, PurePosixPath]: + groups = {} + for membership in memberships: + fields = membership.split(":", 2) + if len(fields) == 3: + for controller in fields[1].split(","): + groups[controller] = PurePosixPath(fields[2]) + return groups + + +def _unescape_mount_field(field: str) -> Path: + return Path(re.sub(r"\\([0-7]{3})", lambda m: chr(int(m[1], 8)), field)) + + +def _cpu_mount( + mount: str, groups: dict[str, PurePosixPath] +) -> tuple[Path, Path, bool] | None: + """Return (cgroup directory, mountpoint, unified) for a CPU cgroup mount.""" + before, separator, after = mount.partition(" - ") + fields, filesystem = before.split(), after.split() + if not separator or len(fields) < 5 or len(filesystem) < 3: + return None + unified = filesystem[0] == "cgroup2" + if not unified and not ( + filesystem[0] == "cgroup" and "cpu" in filesystem[2].split(",") + ): + return None + group = groups.get("" if unified else "cpu") + if group is None: + return None + root, mountpoint = (_unescape_mount_field(field) for field in fields[3:5]) + try: + relative = ( + group.relative_to(root) + if group != PurePosixPath("/") + else PurePosixPath(".") + ) + except ValueError: + return None + if ".." in relative.parts: + return None + return mountpoint / relative, mountpoint, unified + + +def _quota_cpu_count(directory: Path, unified: bool) -> int | None: + try: + if unified: + quota_text, period_text = (directory / "cpu.max").read_text().split() + quota, period = int(quota_text), int(period_text) + else: + quota = int((directory / "cpu.cfs_quota_us").read_text()) + period = int((directory / "cpu.cfs_period_us").read_text()) + except (OSError, ValueError): + return None + if quota > 0 and period > 0: + return (quota + period - 1) // period + return None + + +def _cgroup_cpu_count() -> int | None: + if sys.platform != "linux": + return None + try: + memberships = Path("/proc/self/cgroup").read_text().splitlines() + mounts = Path("/proc/self/mountinfo").read_text().splitlines() + except OSError: + return None + + groups = _cgroup_memberships(memberships) + limits = [] + for mount in mounts: + cpu_mount = _cpu_mount(mount, groups) + if cpu_mount is None: + continue + directory, mountpoint, unified = cpu_mount + while True: + limit = _quota_cpu_count(directory, unified) + if limit is not None: + limits.append(limit) + if directory == mountpoint: + break + directory = directory.parent + return min(limits) if limits else None + + +def available_cpu_count() -> int: + """Respect affinity and CPU time quotas, including visible parent cgroups.""" + count = os.cpu_count() or 1 + try: + count = min(count, len(os.sched_getaffinity(0))) + except (AttributeError, OSError): + pass + quota_count = _cgroup_cpu_count() + if quota_count is not None: + count = min(count, quota_count) + return max(1, count)