[ROCM][NFC] GpuBlaslt matmul thunk cache refactoring part II #23315
Conversation
| absl::StatusOr<se::gpu::BlasLt::MatmulPlan *> | ||
| GetOrCreate(const std::string& key, Func&& create) { |
There was a problem hiding this comment.
This API is dangerous because it's really easy to poison the cache.
The problem is that the user of this function can't rely on the fact that the returned MatmulPlan is the same as what create() would have produced. It might be true right now, but it's brittle. If someone changed the implementation of create and we run two instance of XLA in the same process (which some users do), then this can lead to very subtle bugs.
There was a problem hiding this comment.
I'm wondering whether it's too expensive to always create a MatmulPlan and then keep a list of unique plans in the cache. This would avoid the need for a key as well.
There was a problem hiding this comment.
I mean the idea was here that only GpuBlasLtMatmulThunk can access this cache and that the cache keys are generated using AutotuneCacheKey() constructor: i.e., the one we also for cache keys during GEMM tuning (gemm_algorithm_picker.cc). So I assumed that it is safe to use it in this context: otherwise our autotune algorithm ID would not match to a MatmulPlan which would cause a runtime error.
Yes, potentially one can "precreate" this cache during XLA compilation by going through all computations and extracting all cublas$lt$matmul custom calls (e.g. in some extra HLO pass). But we would have to do this for all GpuExecutors created by XLA since MatmulPlanCache is bound to a particular device.
Or maybe I do not quite understand what you mean by always creating a MatmulPlan ?
| auto& res = meta[device_id]; | ||
| if (!res) res.reset(new MatmulPlanCache()); |
There was a problem hiding this comment.
What about different sessions/contexts on the same GPU? They can't share the plan as far as I know, so this might break things for users that run multiple instances of XLA in the same process.
There was a problem hiding this comment.
I mean previously we have used Stream pointer as a cache key. Each stream is bound to a particular StreamExecutor which I assume is unique for one device (but we can also use a pointer to StreamExecutor as a meta-cache key instead of a device ordinal to make it more safer). MatmulPlans are context-agnostic - the context activation occurs when fetching the algorithms and before the actual Matmul.
There was a problem hiding this comment.
To make it more safe, I have changed the cache interface to use a pointer to cache executor as a key. Therefore, if a user runs several XLA instances within the same process, each of them will create a different StreamExecutor, and hence MatmulPlan caches won't be shared even if they point to the same device ordinal. Though, I assume, most likely, each XLA instance will be bound to a different device (which sounds like a more natural setting).
pemeliya
force-pushed
the
ci_gpublas_lt_cache_refactor2
branch
2 times, most recently
from
eb93eb1 to
28f50f8
Compare
pemeliya
force-pushed
the
ci_gpublas_lt_cache_refactor2
branch
from
28f50f8 to
49ee5c1
Compare
pemeliya
force-pushed
the
ci_gpublas_lt_cache_refactor2
branch
from
49ee5c1 to
1997f00
Compare
|
@beckerhe, I wonder what needs to be done more for this PR ? Is it still pending internal review ? As I wrote, I wrote extra unit tests for gpublaslt_matmul_thunk and matmul_plan_cache |
adapted cublaslt added test enabled tests on rocm update restored command_buffer thunk tests (to be updated in a new PR) moved matmul plan cache to a separate file updated matmul_thunk test updated thunk test added a separate matmul_plan cache unit test made SetAlgorithm function non-const
pemeliya
force-pushed
the
ci_gpublas_lt_cache_refactor2
branch
from
1997f00 to
6ca1c6d
Compare
After this PR is merged: #21886,
GpuBlasLt MatmulPlan cache can now be refactored.
It was originally introduced here: #6595
Here the idea is that multiple matmul thunks can share the same matmul plans (allocated on the same device).
This can significantly reduce the memory overhead for large training. Furthermore, the original cache used stream pointer as a cache key: this might be inefficient when banchmarking the same HLO using XLA tools like multi_host_hlo_runner (which could allocate a new stream for each iteration).
I have also added the correspnding gpublas_lt_matmul_thunk_test to check this functionality.
Besides, I also refactored CublasLtCmd which was a blank copy-paste of CublasLtMatmulThunk.
Finally, I have removed the magic constant 128 and replaced it with GemmConfig::kNumAlgorithms.
@xla-rotation could you please have a look ?
I have also gathered some stats for LLAMA Maxtext model training with 8 GPUs:
So, XLA runtime created 1039 GpuBlasLt Thunk instances while cache contained only 29 unique entries (because most of Gemm configs are duplicates). Hence, provided that it runs with 8 GPUs, we create at least (1039 - 29) * 8 = 8080 matmul plans less than before.