Expanded sharded support for alternative sharding mechanisms #680
+358
−162
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Single-logical-multi-physical sharding allows tensor access between different devices and tighter synchronization on execution. This means that sharding needs to support more than differing device ordinals but also configre multiple queues for the same device. Sharded tensor types are reworked to support tracking both the supported device AND the queue it is enqueued on. To support this each sharded tensor now tracks the DeviceAffinity it is associated with, along with reassigning affinities post construction. This allows pre-sharded models to have their affinities updated with an alternative transfer mechanism. If device affinity is not specified the default arrangement assumes separate device ordinals for each shard.
rsuderman
force-pushed
the
queue_affinities
branch
from
e486ad4 to
9a061e6
Compare
| @@ -279,16 +283,15 @@ def main(): | |||
| tensor_parallelism_size=args.tensor_parallelism_size, | |||
| fake_quant=args.fake_quant, | |||
| ) | |||
| if config.tensor_parallelism_size > 1: | |||
| dataset.root_theta = shard_theta(dataset.root_theta, config) | |||
There was a problem hiding this comment.
Can remove shard_theta import if unused.
|
I saw the following error this morning when attempting to validate toy_llama_tp2 from iree-test-suites by exporting, and compiling with intent to then verify with
python -m sharktank.examples.export_paged_llm_v1 --bs=1 --irpa-file assets/toy_llama_tp2.irpa --output-mlir=llama.mlir --output-config=config.json --use-queue-affinities
iree-compile llama.mlir -o llama.vmfb --iree-hip-target=gfx942 --iree-hal-target-device=hip[0]Received the following error:/toy_new/llama.mlir:4027:12: error: op affinity #hal.device.affinity<@__device_0> is not compatible with the partition affinity #hal.device.affinity<@__device_0, [0]>
%153 = torch.prims.convert_element_type %1, %int5_87 : !torch.vtensor<[256,256],f32>, !torch.int -> !torch.vtensor<[256,256],f16>
^
./toy_new/llama.mlir:4027:12: note: see current operation: %190 = "stream.async.transfer"(%189, %10, %10) <{result_affinity = #hal.device.affinity<@__device_0>, source_affinity = #hal.device.affinity<@__device_0, [1]>}> : (!stream.resource<constant>, index, index) -> !stream.resource<constant>Feedback from Rob this morning before sync: |
| self.rope_dimension_count = rope_dimension_count | ||
| self.max_seqlen = max_seqlen | ||
| self.use_hf = use_hf | ||
| self.static_tables = static_tables | ||
| self.use_table = use_table | ||
|
|
||
| self.rope_freq_base = rope_freq_base if rope_freq_base is not None else 10000.0 | ||
| self.tensor_parallelism_size = tensor_parallelism_size | ||
| self.devices = devices |
There was a problem hiding this comment.
This is redundant with L34, can be removed.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Single-logical-multi-physical sharding allows tensor access between
different devices and tighter synchronization on execution. This means
that sharding needs to support more than differing device ordinals but
also configre multiple queues for the same device. Sharded tensor types
are reworked to support tracking both the supported device AND the queue
it is enqueued on.
To support this each sharded tensor now tracks the DeviceAffinity it is
associated with, along with reassigning affinities post construction.
This allows pre-sharded models to have their affinities updated with an
alternative transfer mechanism.
If device affinity is not specified the default arrangement assumes
separate device ordinals for each shard.