Lunderberg commented on PR #15700: URL: https://github.com/apache/tvm/pull/15700#issuecomment-1711894095
Regarding disco, I agree that the distribution will be handled through the disco runtime, but don't see any conflict with that and the `LiftTransformParams` approach. By using both at the same time, we can have a simpler and more efficient startup. ```python # Load weights, then use disco runtime to perform sharding. The # sharding function uses the rank_dref to determine which portion of # each weight to retain. rank_dref = ... weights = tvm.nd.array(...) sharding_function = disco_session.get_func(...) sharded_weights_dref = sharding_function(weights, rank_dref) # Load sharded weights through disco runtime. The load weights # function uses the rank_dref to determine which set of pre-sharded # weights to load. rank_dref = ... load_weights_function = disco_session.get_func(...) sharded_weights_dref = load_weights_function(rank_dref) ``` Rather than loading all the weights onto one GPU, then sharding them at runtime, using both tools together allows us to have the pre-sharded weights loaded only onto the GPU that will be using that shard. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
