Skip to content

Commit 0da8693

Browse files
MarkDaoustcopybara-github
MarkDaoust
authored and
copybara-github
committed
DTensor docs update.
* Simplify data-handling in the keras-tutorial. * Add a couple more figures. PiperOrigin-RevId: 446486657
1 parent 279f9df commit 0da8693

File tree

2 files changed

+201
-163
lines changed

2 files changed

+201
-163
lines changed

‎site/en/guide/dtensor_overview.ipynb

Lines changed: 67 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@
174174
"source": [
175175
"In a 1 dimensional `Mesh`, all devices form a list in a single mesh dimension. The following example uses `dtensor.create_mesh` to create a mesh from 6 CPU devices along a mesh dimension `'x'` with a size of 6 devices:\n",
176176
"\n",
177-
"<img src=\"https://www.tensorflow.org/guide/images/dtensor_mesh_1d.png\" alt=\"A 1 dimensional mesh with 6 CPUs\" class=\"no-filter\">\n"
177+
"<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_mesh_1d.png\" alt=\"A 1 dimensional mesh with 6 CPUs\" class=\"no-filter\">\n"
178178
]
179179
},
180180
{
@@ -197,7 +197,7 @@
197197
"source": [
198198
"A `Mesh` can be multi dimensional as well. In the following example, 6 CPU devices form a `3x2` mesh, where the `'x'` mesh dimension has a size of 3 devices, and the `'y'` mesh dimension has a size of 2 devices:\n",
199199
"\n",
200-
"<img src=\"https://www.tensorflow.org/guide/images/dtensor_mesh_2d.png\" alt=\"A 2 dimensional mesh with 6 CPUs\"\n",
200+
"<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_mesh_2d.png\" alt=\"A 2 dimensional mesh with 6 CPUs\"\n",
201201
" class=\"no-filter\">"
202202
]
203203
},
@@ -221,15 +221,16 @@
221221
"source": [
222222
"### Layout\n",
223223
"\n",
224-
"**`Layout`** specifies how a tensor is distributed, or sharded, on a `Mesh`. In DTensor, the placement specification is on a per-axis bases. An axis of a `Layout` can be either `sharded` or `unsharded` (replicated) along a mesh dimension.\n",
224+
"**`Layout`** specifies how a tensor is distributed, or sharded, on a `Mesh`.\n",
225225
"\n",
226226
"Note: In order to avoid confusions between `Mesh` and `Layout`, the term *dimension* is always associated with `Mesh`, and the term *axis* with `Tensor` and `Layout` in this guide.\n",
227227
"\n",
228-
"The rank of a `Layout` and the number of dimensions of a `Mesh` do not need to match. The `unsharded` axes of a `Layout` do not need to be associated to a mesh dimension, and `unsharded` mesh dimensions do not need to be associated with a `layout` axis.\n",
228+
"The rank of `Layout` should be the same as the rank of the `Tensor` where the `Layout` is applied. For each of the `Tensor`'s axes the `Layout` may specifiy a mesh dimension to shard the tensor across, or specify the axis as \"unsharded\".\n",
229+
"The tensor is replicated across any mesh dimensions that it is not sharded across.\n",
229230
"\n",
230-
"On the other hand, the rank of `Layout` should be the same as the rank of the `Tensor` where the `Layout` is applied.\n",
231+
"The rank of a `Layout` and the number of dimensions of a `Mesh` do not need to match. The `unsharded` axes of a `Layout` do not need to be associated to a mesh dimension, and `unsharded` mesh dimensions do not need to be associated with a `layout` axis.\n",
231232
"\n",
232-
"<img src=\"https://www.tensorflow.org/guide/images/dtensor_components_diag.png\" alt=\"Diagram of dtensor components.\"\n",
233+
"<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_components_diag.png\" alt=\"Diagram of dtensor components.\"\n",
233234
" class=\"no-filter\">"
234235
]
235236
},
@@ -248,9 +249,8 @@
248249
"id": "fqzCNlWAbm-c"
249250
},
250251
"source": [
251-
"On a 1-dimensional mesh such as `[(\"x\", 6)]` (`mesh_1d` in the previous section), `Layout([\"unsharded\"], mesh_1d)` is a layout for a rank-1 tensor replicated on 6 devices.\n",
252-
"\n",
253-
"<img src=\"https://www.tensorflow.org/guide/images/dtensor_layout_rank1.png\" alt=\"Layout for a rank-1 tensor\" class=\"no-filter\">"
252+
"On a 1-dimensional mesh such as `[(\"x\", 6)]` (`mesh_1d` in the previous section), `Layout([\"unsharded\", \"unsharded\"], mesh_1d)` is a layout for a rank-2 tensor replicated across 6 devices.\n",
253+
"<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_layout_replicated.png\" alt=\"A tensor replicated across a rank-1 mesh\" class=\"no-filter\">"
254254
]
255255
},
256256
{
@@ -261,19 +261,47 @@
261261
},
262262
"outputs": [],
263263
"source": [
264-
"layout_rank_1_1d = dtensor.Layout([dtensor.UNSHARDED], mesh_1d)\n",
265-
"print(layout_rank_1_1d)"
264+
"layout = dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED], mesh_1d)"
266265
]
267266
},
268267
{
269268
"cell_type": "markdown",
270269
"metadata": {
271-
"id": "1Kyg0V3ehMNJ"
270+
"id": "ywRJwuLDt2Qq"
272271
},
273272
"source": [
274-
"Given a 2-dimensional 3x2 mesh such as `[(\"x\", 3), (\"y\", 2)]`, (`mesh_2d` from the previous section), `Layout([\"x\", dtensor.UNSHARDED], mesh_2d)` is a layout for a rank-2 `Tensor`, whose first axis is sharded on mesh dimension `x`, and second axis is replicated (`UNSHARDED`).\n",
273+
"Using the same tensor and mesh the layout `Layout(['unsharded', 'x'])` would shard the second axis of the tensor across the 6 devices.\n",
275274
"\n",
276-
"<img src=\"https://www.tensorflow.org/guide/images/dtensor_layout_rank2.png\" alt=\"Layout for a rank-2 tensor\" class=\"no-filter\">\n"
275+
"<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_layout_rank1.png\" alt=\"A tensor sharded across a rank-1 mesh\" class=\"no-filter\">"
276+
]
277+
},
278+
{
279+
"cell_type": "code",
280+
"execution_count": null,
281+
"metadata": {
282+
"id": "7BgqL0jUvV5a"
283+
},
284+
"outputs": [],
285+
"source": [
286+
"layout = dtensor.Layout([dtensor.UNSHARDED, 'x'], mesh_1d)"
287+
]
288+
},
289+
{
290+
"cell_type": "markdown",
291+
"metadata": {
292+
"id": "DgciDNmK76l9"
293+
},
294+
"source": [
295+
"Given a 2-dimensional 3x2 mesh such as `[(\"x\", 3), (\"y\", 2)]`, (`mesh_2d` from the previous section), `Layout([\"y\", \"x\"], mesh_2d)` is a layout for a rank-2 `Tensor` whose first axis is sharded across across mesh dimension `\"y\"`, and whose second axis is sharded across mesh dimension `\"x\"`."
296+
]
297+
},
298+
{
299+
"cell_type": "markdown",
300+
"metadata": {
301+
"id": "Eyp_qOSyvieo"
302+
},
303+
"source": [
304+
"<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_layout_rank2.png\" alt=\"A tensorr with it's first axis sharded across mesh dimension 'y' and it's second axis sharded across mesh dimension 'x'\" class=\"no-filter\">\n"
277305
]
278306
},
279307
{
@@ -284,8 +312,29 @@
284312
},
285313
"outputs": [],
286314
"source": [
287-
"layout_rank_2_2d = dtensor.Layout(['x', dtensor.UNSHARDED], mesh_2d)\n",
288-
"print(layout_rank_2_2d)"
315+
"layout = dtensor.Layout(['y', 'x'], mesh_2d)"
316+
]
317+
},
318+
{
319+
"cell_type": "markdown",
320+
"metadata": {
321+
"id": "1Kyg0V3ehMNJ"
322+
},
323+
"source": [
324+
"For the same `mesh_2d`, the layout `Layout([\"x\", dtensor.UNSHARDED], mesh_2d)` is a layout for a rank-2 `Tensor` that is replicated across `\"y\"`, and whose first axis is sharded on mesh dimension `x`.\n",
325+
"\n",
326+
"<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_layout_hybrid.png\" alt=\"A tensor replicated across mesh-dimension y, with it's first axis sharded across mesh dimension 'x'\" class=\"no-filter\">\n"
327+
]
328+
},
329+
{
330+
"cell_type": "code",
331+
"execution_count": null,
332+
"metadata": {
333+
"id": "IkWe6mVl7uRb"
334+
},
335+
"outputs": [],
336+
"source": [
337+
"layout = dtensor.Layout([\"x\", dtensor.UNSHARDED], mesh_2d)"
289338
]
290339
},
291340
{
@@ -450,8 +499,7 @@
450499
"source": [
451500
"The inverse operation of `dtensor.unpack` is `dtensor.pack`. Component tensors can be packed back into a DTensor.\n",
452501
"\n",
453-
"The components must have the same rank and dtype, which will be the rank and dtype of the returned DTensor. However there is no strict requirement on the device placement of component tensors as inputs of `dtensor.unpack`: the function will automatically copy the component tensors to their respective corresponding devices.\n",
454-
"\n"
502+
"The components must have the same rank and dtype, which will be the rank and dtype of the returned DTensor. However there is no strict requirement on the device placement of component tensors as inputs of `dtensor.unpack`: the function will automatically copy the component tensors to their respective corresponding devices.\n"
455503
]
456504
},
457505
{
@@ -599,7 +647,7 @@
599647
"source": [
600648
"You can inspect the component tensors of the created DTensor and verify they are indeed sharded according to your scheme. It may be helpful to illustrate the situation with a chart:\n",
601649
"\n",
602-
" <img src=\"https://www.tensorflow.org/guide/images/dtensor_hybrid_mesh.png\" alt=\"A 3x2 hybrid mesh with 6 CPUs\"\n",
650+
" <img src=\"https://www.tensorflow.org/images/dtensor/dtensor_hybrid_mesh.png\" alt=\"A 3x2 hybrid mesh with 6 CPUs\"\n",
603651
" class=\"no-filter\" width=75%>\n"
604652
]
605653
},

0 commit comments

Comments
 (0)