We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] Horovod example scripts fail when user supplies batch size parameter
**Describe the bug**
Using the batch size parameter on the TF Horovod example causes a type error with a mismatch between str and int.

**Steps/Code to reproduce bug**
Run the TF Horovod example with the arguments `--b_size 1024`.

**Expected behavior**
The script should accept a user-provided batch size.

**Environment details (please complete the following information):**
 - Environment location: Bare-metal
 - Method of NVTabular install: conda
 
**Additional context**
I believe [this line](https://github.com/NVIDIA/NVTabular/blob/main/examples/horovod/tf_hvd_simple.py#L30) and the same line in the Torch example just need type coercions from str to int.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/horovod/tf_hvd_simple.py`
Content:
```
1 # External dependencies
2 import argparse
3 import glob
4 import os
5 
6 import cupy
7 
8 # we can control how much memory to give tensorflow with this environment variable
9 # IMPORTANT: make sure you do this before you initialize TF's runtime, otherwise
10 # TF will have claimed all free GPU memory
11 os.environ["TF_MEMORY_ALLOCATION"] = "0.3"  # fraction of free memory
12 import horovod.tensorflow as hvd  # noqa: E402
13 import tensorflow as tf  # noqa: E402
14 
15 import nvtabular as nvt  # noqa: E402
16 from nvtabular.framework_utils.tensorflow import layers  # noqa: E402
17 from nvtabular.loader.tensorflow import KerasSequenceLoader  # noqa: E402
18 
19 parser = argparse.ArgumentParser(description="Process some integers.")
20 parser.add_argument("--dir_in", default=None, help="Input directory")
21 parser.add_argument("--b_size", default=None, help="batch size")
22 parser.add_argument("--cats", default=None, help="categorical columns")
23 parser.add_argument("--cats_mh", default=None, help="categorical multihot columns")
24 parser.add_argument("--conts", default=None, help="continuous columns")
25 parser.add_argument("--labels", default=None, help="continuous columns")
26 args = parser.parse_args()
27 
28 
29 BASE_DIR = args.dir_in or "./data/"
30 BATCH_SIZE = args.b_size or 16384  # Batch Size
31 CATEGORICAL_COLUMNS = args.cats or ["movieId", "userId"]  # Single-hot
32 CATEGORICAL_MH_COLUMNS = args.cats_mh or ["genres"]  # Multi-hot
33 NUMERIC_COLUMNS = args.conts or []
34 TRAIN_PATHS = sorted(
35     glob.glob(os.path.join(BASE_DIR, "train/*.parquet"))
36 )  # Output from ETL-with-NVTabular
37 hvd.init()
38 
39 # Seed with system randomness (or a static seed)
40 cupy.random.seed(None)
41 
42 
43 def seed_fn():
44     """
45     Generate consistent dataloader shuffle seeds across workers
46 
47     Reseeds each worker's dataloader each epoch to get fresh a shuffle
48     that's consistent across workers.
49     """
50     min_int, max_int = tf.int32.limits
51     max_rand = max_int // hvd.size()
52 
53     # Generate a seed fragment
54     seed_fragment = cupy.random.randint(0, max_rand).get()
55 
56     # Aggregate seed fragments from all Horovod workers
57     seed_tensor = tf.constant(seed_fragment)
58     reduced_seed = hvd.allreduce(seed_tensor, name="shuffle_seed", op=hvd.mpi_ops.Sum) % max_rand
59 
60     return reduced_seed
61 
62 
63 proc = nvt.Workflow.load(os.path.join(BASE_DIR, "workflow/"))
64 EMBEDDING_TABLE_SHAPES = nvt.ops.get_embedding_sizes(proc)
65 
66 train_dataset_tf = KerasSequenceLoader(
67     TRAIN_PATHS,  # you could also use a glob pattern
68     batch_size=BATCH_SIZE,
69     label_names=["rating"],
70     cat_names=CATEGORICAL_COLUMNS + CATEGORICAL_MH_COLUMNS,
71     cont_names=NUMERIC_COLUMNS,
72     engine="parquet",
73     shuffle=True,
74     seed_fn=seed_fn,
75     buffer_size=0.06,  # how many batches to load at once
76     parts_per_chunk=1,
77     global_size=hvd.size(),
78     global_rank=hvd.rank(),
79 )
80 inputs = {}  # tf.keras.Input placeholders for each feature to be used
81 emb_layers = []  # output of all embedding layers, which will be concatenated
82 for col in CATEGORICAL_COLUMNS:
83     inputs[col] = tf.keras.Input(name=col, dtype=tf.int32, shape=(1,))
84 # Note that we need two input tensors for multi-hot categorical features
85 for col in CATEGORICAL_MH_COLUMNS:
86     inputs[col + "__values"] = tf.keras.Input(name=f"{col}__values", dtype=tf.int64, shape=(1,))
87     inputs[col + "__nnzs"] = tf.keras.Input(name=f"{col}__nnzs", dtype=tf.int64, shape=(1,))
88 for col in CATEGORICAL_COLUMNS + CATEGORICAL_MH_COLUMNS:
89     emb_layers.append(
90         tf.feature_column.embedding_column(
91             tf.feature_column.categorical_column_with_identity(
92                 col, EMBEDDING_TABLE_SHAPES[col][0]  # Input dimension (vocab size)
93             ),
94             EMBEDDING_TABLE_SHAPES[col][1],  # Embedding output dimension
95         )
96     )
97 emb_layer = layers.DenseFeatures(emb_layers)
98 x_emb_output = emb_layer(inputs)
99 x = tf.keras.layers.Dense(128, activation="relu")(x_emb_output)
100 x = tf.keras.layers.Dense(128, activation="relu")(x)
101 x = tf.keras.layers.Dense(128, activation="relu")(x)
102 x = tf.keras.layers.Dense(1, activation="sigmoid")(x)
103 model = tf.keras.Model(inputs=inputs, outputs=x)
104 loss = tf.losses.BinaryCrossentropy()
105 opt = tf.keras.optimizers.SGD(0.01 * hvd.size())
106 opt = hvd.DistributedOptimizer(opt)
107 checkpoint_dir = "./checkpoints"
108 checkpoint = tf.train.Checkpoint(model=model, optimizer=opt)
109 
110 
111 @tf.function(experimental_relax_shapes=True)
112 def training_step(examples, labels, first_batch):
113     with tf.GradientTape() as tape:
114         probs = model(examples, training=True)
115         loss_value = loss(labels, probs)
116     # Horovod: add Horovod Distributed GradientTape.
117     tape = hvd.DistributedGradientTape(tape, sparse_as_dense=True)
118     grads = tape.gradient(loss_value, model.trainable_variables)
119     opt.apply_gradients(zip(grads, model.trainable_variables))
120     # Horovod: broadcast initial variable states from rank 0 to all other processes.
121     # This is necessary to ensure consistent initialization of all workers when
122     # training is started with random weights or restored from a checkpoint.
123     #
124     # Note: broadcast should be done after the first gradient step to ensure optimizer
125     # initialization.
126     if first_batch:
127         hvd.broadcast_variables(model.variables, root_rank=0)
128         hvd.broadcast_variables(opt.variables(), root_rank=0)
129     return loss_value
130 
131 
132 # Horovod: adjust number of steps based on number of GPUs.
133 for batch, (examples, labels) in enumerate(train_dataset_tf):
134     loss_value = training_step(examples, labels, batch == 0)
135     if batch % 10 == 0 and hvd.local_rank() == 0:
136         print("Step #%d\tLoss: %.6f" % (batch, loss_value))
137 hvd.join()
138 # Horovod: save checkpoints only on worker 0 to prevent other workers from
139 # corrupting it.
140 if hvd.rank() == 0:
141     checkpoint.save(checkpoint_dir)
142
```

Path: `examples/horovod/torch-nvt-horovod.py`
Content:
```
1 import argparse
2 import glob
3 import os
4 from time import time
5 
6 import cupy
7 import torch
8 
9 import nvtabular as nvt
10 from nvtabular.framework_utils.torch.models import Model
11 from nvtabular.framework_utils.torch.utils import process_epoch
12 from nvtabular.loader.torch import DLDataLoader, TorchAsyncItr
13 
14 # Horovod must be the last import to avoid conflicts
15 import horovod.torch as hvd  # noqa: E402, isort:skip
16 
17 
18 parser = argparse.ArgumentParser(description="Train a multi-gpu model with Torch and Horovod")
19 parser.add_argument("--dir_in", default=None, help="Input directory")
20 parser.add_argument("--batch_size", default=None, help="Batch size")
21 parser.add_argument("--cats", default=None, help="Categorical columns")
22 parser.add_argument("--cats_mh", default=None, help="Categorical multihot columns")
23 parser.add_argument("--conts", default=None, help="Continuous columns")
24 parser.add_argument("--labels", default=None, help="Label columns")
25 parser.add_argument("--epochs", default=1, help="Training epochs")
26 args = parser.parse_args()
27 
28 hvd.init()
29 
30 gpu_to_use = hvd.local_rank()
31 
32 if torch.cuda.is_available():
33     torch.cuda.set_device(gpu_to_use)
34 
35 
36 BASE_DIR = os.path.expanduser(args.dir_in or "./data/")
37 BATCH_SIZE = args.batch_size or 16384  # Batch Size
38 CATEGORICAL_COLUMNS = args.cats or ["movieId", "userId"]  # Single-hot
39 CATEGORICAL_MH_COLUMNS = args.cats_mh or ["genres"]  # Multi-hot
40 NUMERIC_COLUMNS = args.conts or []
41 
42 # Output from ETL-with-NVTabular
43 TRAIN_PATHS = sorted(glob.glob(os.path.join(BASE_DIR, "train", "*.parquet")))
44 
45 proc = nvt.Workflow.load(os.path.join(BASE_DIR, "workflow/"))
46 
47 EMBEDDING_TABLE_SHAPES = nvt.ops.get_embedding_sizes(proc)
48 
49 
50 # TensorItrDataset returns a single batch of x_cat, x_cont, y.
51 def collate_fn(x):
52     return x
53 
54 
55 # Seed with system randomness (or a static seed)
56 cupy.random.seed(None)
57 
58 
59 def seed_fn():
60     """
61     Generate consistent dataloader shuffle seeds across workers
62 
63     Reseeds each worker's dataloader each epoch to get fresh a shuffle
64     that's consistent across workers.
65     """
66 
67     max_rand = torch.iinfo(torch.int).max // hvd.size()
68 
69     # Generate a seed fragment
70     seed_fragment = cupy.random.randint(0, max_rand)
71 
72     # Aggregate seed fragments from all Horovod workers
73     seed_tensor = torch.tensor(seed_fragment)
74     reduced_seed = hvd.allreduce(seed_tensor, name="shuffle_seed", op=hvd.mpi_ops.Sum) % max_rand
75 
76     return reduced_seed
77 
78 
79 train_dataset = TorchAsyncItr(
80     nvt.Dataset(TRAIN_PATHS),
81     batch_size=BATCH_SIZE,
82     cats=CATEGORICAL_COLUMNS + CATEGORICAL_MH_COLUMNS,
83     conts=NUMERIC_COLUMNS,
84     labels=["rating"],
85     devices=[gpu_to_use],
86     global_size=hvd.size(),
87     global_rank=hvd.rank(),
88     shuffle=True,
89     seed_fn=seed_fn,
90 )
91 train_loader = DLDataLoader(
92     train_dataset, batch_size=None, collate_fn=collate_fn, pin_memory=False, num_workers=0
93 )
94 
95 
96 EMBEDDING_TABLE_SHAPES_TUPLE = (
97     {
98         CATEGORICAL_COLUMNS[0]: EMBEDDING_TABLE_SHAPES[CATEGORICAL_COLUMNS[0]],
99         CATEGORICAL_COLUMNS[1]: EMBEDDING_TABLE_SHAPES[CATEGORICAL_COLUMNS[1]],
100     },
101     {CATEGORICAL_MH_COLUMNS[0]: EMBEDDING_TABLE_SHAPES[CATEGORICAL_MH_COLUMNS[0]]},
102 )
103 
104 model = Model(
105     embedding_table_shapes=EMBEDDING_TABLE_SHAPES_TUPLE,
106     num_continuous=0,
107     emb_dropout=0.0,
108     layer_hidden_dims=[128, 128, 128],
109     layer_dropout_rates=[0.0, 0.0, 0.0],
110 ).cuda()
111 
112 lr_scaler = hvd.size()
113 
114 optimizer = torch.optim.Adam(model.parameters(), lr=0.01 * lr_scaler)
115 
116 hvd.broadcast_parameters(model.state_dict(), root_rank=0)
117 hvd.broadcast_optimizer_state(optimizer, root_rank=0)
118 
119 optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
120 
121 for epoch in range(args.epochs):
122     start = time()
123     print(f"Training epoch {epoch}")
124     train_loss, y_pred, y = process_epoch(train_loader, model, train=True, optimizer=optimizer)
125     hvd.join(gpu_to_use)
126     hvd.broadcast_parameters(model.state_dict(), root_rank=0)
127     print(f"Epoch {epoch:02d}. Train loss: {train_loss:.4f}.")
128     hvd.join(gpu_to_use)
129     t_final = time() - start
130     total_rows = train_dataset.num_rows_processed
131     print(
132         f"run_time: {t_final} - rows: {total_rows} - "
133         f"epochs: {epoch} - dl_thru: {total_rows / t_final}"
134     )
135 
136 
137 hvd.join(gpu_to_use)
138 if hvd.local_rank() == 0:
139     print("Training complete")
140
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/horovod/tf_hvd_simple.py b/examples/horovod/tf_hvd_simple.py
--- a/examples/horovod/tf_hvd_simple.py
+++ b/examples/horovod/tf_hvd_simple.py
@@ -18,7 +18,7 @@
 
 parser = argparse.ArgumentParser(description="Process some integers.")
 parser.add_argument("--dir_in", default=None, help="Input directory")
-parser.add_argument("--b_size", default=None, help="batch size")
+parser.add_argument("--batch_size", default=None, help="batch size")
 parser.add_argument("--cats", default=None, help="categorical columns")
 parser.add_argument("--cats_mh", default=None, help="categorical multihot columns")
 parser.add_argument("--conts", default=None, help="continuous columns")
@@ -27,7 +27,7 @@
 
 
 BASE_DIR = args.dir_in or "./data/"
-BATCH_SIZE = args.b_size or 16384  # Batch Size
+BATCH_SIZE = int(args.batch_size) or 16384  # Batch Size
 CATEGORICAL_COLUMNS = args.cats or ["movieId", "userId"]  # Single-hot
 CATEGORICAL_MH_COLUMNS = args.cats_mh or ["genres"]  # Multi-hot
 NUMERIC_COLUMNS = args.conts or []
diff --git a/examples/horovod/torch-nvt-horovod.py b/examples/horovod/torch-nvt-horovod.py
--- a/examples/horovod/torch-nvt-horovod.py
+++ b/examples/horovod/torch-nvt-horovod.py
@@ -34,7 +34,7 @@
 
 
 BASE_DIR = os.path.expanduser(args.dir_in or "./data/")
-BATCH_SIZE = args.batch_size or 16384  # Batch Size
+BATCH_SIZE = int(args.batch_size) or 16384  # Batch Size
 CATEGORICAL_COLUMNS = args.cats or ["movieId", "userId"]  # Single-hot
 CATEGORICAL_MH_COLUMNS = args.cats_mh or ["genres"]  # Multi-hot
 NUMERIC_COLUMNS = args.conts or []

{"golden_diff": "diff --git a/examples/horovod/tf_hvd_simple.py b/examples/horovod/tf_hvd_simple.py\n--- a/examples/horovod/tf_hvd_simple.py\n+++ b/examples/horovod/tf_hvd_simple.py\n@@ -18,7 +18,7 @@\n \n parser = argparse.ArgumentParser(description=\"Process some integers.\")\n parser.add_argument(\"--dir_in\", default=None, help=\"Input directory\")\n-parser.add_argument(\"--b_size\", default=None, help=\"batch size\")\n+parser.add_argument(\"--batch_size\", default=None, help=\"batch size\")\n parser.add_argument(\"--cats\", default=None, help=\"categorical columns\")\n parser.add_argument(\"--cats_mh\", default=None, help=\"categorical multihot columns\")\n parser.add_argument(\"--conts\", default=None, help=\"continuous columns\")\n@@ -27,7 +27,7 @@\n \n \n BASE_DIR = args.dir_in or \"./data/\"\n-BATCH_SIZE = args.b_size or 16384  # Batch Size\n+BATCH_SIZE = int(args.batch_size) or 16384  # Batch Size\n CATEGORICAL_COLUMNS = args.cats or [\"movieId\", \"userId\"]  # Single-hot\n CATEGORICAL_MH_COLUMNS = args.cats_mh or [\"genres\"]  # Multi-hot\n NUMERIC_COLUMNS = args.conts or []\ndiff --git a/examples/horovod/torch-nvt-horovod.py b/examples/horovod/torch-nvt-horovod.py\n--- a/examples/horovod/torch-nvt-horovod.py\n+++ b/examples/horovod/torch-nvt-horovod.py\n@@ -34,7 +34,7 @@\n \n \n BASE_DIR = os.path.expanduser(args.dir_in or \"./data/\")\n-BATCH_SIZE = args.batch_size or 16384  # Batch Size\n+BATCH_SIZE = int(args.batch_size) or 16384  # Batch Size\n CATEGORICAL_COLUMNS = args.cats or [\"movieId\", \"userId\"]  # Single-hot\n CATEGORICAL_MH_COLUMNS = args.cats_mh or [\"genres\"]  # Multi-hot\n NUMERIC_COLUMNS = args.conts or []\n", "issue": "[BUG] Horovod example scripts fail when user supplies batch size parameter\n**Describe the bug**\r\nUsing the batch size parameter on the TF Horovod example causes a type error with a mismatch between str and int.\r\n\r\n**Steps/Code to reproduce bug**\r\nRun the TF Horovod example with the arguments `--b_size 1024`.\r\n\r\n**Expected behavior**\r\nThe script should accept a user-provided batch size.\r\n\r\n**Environment details (please complete the following information):**\r\n - Environment location: Bare-metal\r\n - Method of NVTabular install: conda\r\n \r\n**Additional context**\r\nI believe [this line](https://github.com/NVIDIA/NVTabular/blob/main/examples/horovod/tf_hvd_simple.py#L30) and the same line in the Torch example just need type coercions from str to int.\r\n\n", "before_files": [{"content": "# External dependencies\nimport argparse\nimport glob\nimport os\n\nimport cupy\n\n# we can control how much memory to give tensorflow with this environment variable\n# IMPORTANT: make sure you do this before you initialize TF's runtime, otherwise\n# TF will have claimed all free GPU memory\nos.environ[\"TF_MEMORY_ALLOCATION\"] = \"0.3\"  # fraction of free memory\nimport horovod.tensorflow as hvd  # noqa: E402\nimport tensorflow as tf  # noqa: E402\n\nimport nvtabular as nvt  # noqa: E402\nfrom nvtabular.framework_utils.tensorflow import layers  # noqa: E402\nfrom nvtabular.loader.tensorflow import KerasSequenceLoader  # noqa: E402\n\nparser = argparse.ArgumentParser(description=\"Process some integers.\")\nparser.add_argument(\"--dir_in\", default=None, help=\"Input directory\")\nparser.add_argument(\"--b_size\", default=None, help=\"batch size\")\nparser.add_argument(\"--cats\", default=None, help=\"categorical columns\")\nparser.add_argument(\"--cats_mh\", default=None, help=\"categorical multihot columns\")\nparser.add_argument(\"--conts\", default=None, help=\"continuous columns\")\nparser.add_argument(\"--labels\", default=None, help=\"continuous columns\")\nargs = parser.parse_args()\n\n\nBASE_DIR = args.dir_in or \"./data/\"\nBATCH_SIZE = args.b_size or 16384  # Batch Size\nCATEGORICAL_COLUMNS = args.cats or [\"movieId\", \"userId\"]  # Single-hot\nCATEGORICAL_MH_COLUMNS = args.cats_mh or [\"genres\"]  # Multi-hot\nNUMERIC_COLUMNS = args.conts or []\nTRAIN_PATHS = sorted(\n    glob.glob(os.path.join(BASE_DIR, \"train/*.parquet\"))\n)  # Output from ETL-with-NVTabular\nhvd.init()\n\n# Seed with system randomness (or a static seed)\ncupy.random.seed(None)\n\n\ndef seed_fn():\n    \"\"\"\n    Generate consistent dataloader shuffle seeds across workers\n\n    Reseeds each worker's dataloader each epoch to get fresh a shuffle\n    that's consistent across workers.\n    \"\"\"\n    min_int, max_int = tf.int32.limits\n    max_rand = max_int // hvd.size()\n\n    # Generate a seed fragment\n    seed_fragment = cupy.random.randint(0, max_rand).get()\n\n    # Aggregate seed fragments from all Horovod workers\n    seed_tensor = tf.constant(seed_fragment)\n    reduced_seed = hvd.allreduce(seed_tensor, name=\"shuffle_seed\", op=hvd.mpi_ops.Sum) % max_rand\n\n    return reduced_seed\n\n\nproc = nvt.Workflow.load(os.path.join(BASE_DIR, \"workflow/\"))\nEMBEDDING_TABLE_SHAPES = nvt.ops.get_embedding_sizes(proc)\n\ntrain_dataset_tf = KerasSequenceLoader(\n    TRAIN_PATHS,  # you could also use a glob pattern\n    batch_size=BATCH_SIZE,\n    label_names=[\"rating\"],\n    cat_names=CATEGORICAL_COLUMNS + CATEGORICAL_MH_COLUMNS,\n    cont_names=NUMERIC_COLUMNS,\n    engine=\"parquet\",\n    shuffle=True,\n    seed_fn=seed_fn,\n    buffer_size=0.06,  # how many batches to load at once\n    parts_per_chunk=1,\n    global_size=hvd.size(),\n    global_rank=hvd.rank(),\n)\ninputs = {}  # tf.keras.Input placeholders for each feature to be used\nemb_layers = []  # output of all embedding layers, which will be concatenated\nfor col in CATEGORICAL_COLUMNS:\n    inputs[col] = tf.keras.Input(name=col, dtype=tf.int32, shape=(1,))\n# Note that we need two input tensors for multi-hot categorical features\nfor col in CATEGORICAL_MH_COLUMNS:\n    inputs[col + \"__values\"] = tf.keras.Input(name=f\"{col}__values\", dtype=tf.int64, shape=(1,))\n    inputs[col + \"__nnzs\"] = tf.keras.Input(name=f\"{col}__nnzs\", dtype=tf.int64, shape=(1,))\nfor col in CATEGORICAL_COLUMNS + CATEGORICAL_MH_COLUMNS:\n    emb_layers.append(\n        tf.feature_column.embedding_column(\n            tf.feature_column.categorical_column_with_identity(\n                col, EMBEDDING_TABLE_SHAPES[col][0]  # Input dimension (vocab size)\n            ),\n            EMBEDDING_TABLE_SHAPES[col][1],  # Embedding output dimension\n        )\n    )\nemb_layer = layers.DenseFeatures(emb_layers)\nx_emb_output = emb_layer(inputs)\nx = tf.keras.layers.Dense(128, activation=\"relu\")(x_emb_output)\nx = tf.keras.layers.Dense(128, activation=\"relu\")(x)\nx = tf.keras.layers.Dense(128, activation=\"relu\")(x)\nx = tf.keras.layers.Dense(1, activation=\"sigmoid\")(x)\nmodel = tf.keras.Model(inputs=inputs, outputs=x)\nloss = tf.losses.BinaryCrossentropy()\nopt = tf.keras.optimizers.SGD(0.01 * hvd.size())\nopt = hvd.DistributedOptimizer(opt)\ncheckpoint_dir = \"./checkpoints\"\ncheckpoint = tf.train.Checkpoint(model=model, optimizer=opt)\n\n\n@tf.function(experimental_relax_shapes=True)\ndef training_step(examples, labels, first_batch):\n    with tf.GradientTape() as tape:\n        probs = model(examples, training=True)\n        loss_value = loss(labels, probs)\n    # Horovod: add Horovod Distributed GradientTape.\n    tape = hvd.DistributedGradientTape(tape, sparse_as_dense=True)\n    grads = tape.gradient(loss_value, model.trainable_variables)\n    opt.apply_gradients(zip(grads, model.trainable_variables))\n    # Horovod: broadcast initial variable states from rank 0 to all other processes.\n    # This is necessary to ensure consistent initialization of all workers when\n    # training is started with random weights or restored from a checkpoint.\n    #\n    # Note: broadcast should be done after the first gradient step to ensure optimizer\n    # initialization.\n    if first_batch:\n        hvd.broadcast_variables(model.variables, root_rank=0)\n        hvd.broadcast_variables(opt.variables(), root_rank=0)\n    return loss_value\n\n\n# Horovod: adjust number of steps based on number of GPUs.\nfor batch, (examples, labels) in enumerate(train_dataset_tf):\n    loss_value = training_step(examples, labels, batch == 0)\n    if batch % 10 == 0 and hvd.local_rank() == 0:\n        print(\"Step #%d\\tLoss: %.6f\" % (batch, loss_value))\nhvd.join()\n# Horovod: save checkpoints only on worker 0 to prevent other workers from\n# corrupting it.\nif hvd.rank() == 0:\n    checkpoint.save(checkpoint_dir)\n", "path": "examples/horovod/tf_hvd_simple.py"}, {"content": "import argparse\nimport glob\nimport os\nfrom time import time\n\nimport cupy\nimport torch\n\nimport nvtabular as nvt\nfrom nvtabular.framework_utils.torch.models import Model\nfrom nvtabular.framework_utils.torch.utils import process_epoch\nfrom nvtabular.loader.torch import DLDataLoader, TorchAsyncItr\n\n# Horovod must be the last import to avoid conflicts\nimport horovod.torch as hvd  # noqa: E402, isort:skip\n\n\nparser = argparse.ArgumentParser(description=\"Train a multi-gpu model with Torch and Horovod\")\nparser.add_argument(\"--dir_in\", default=None, help=\"Input directory\")\nparser.add_argument(\"--batch_size\", default=None, help=\"Batch size\")\nparser.add_argument(\"--cats\", default=None, help=\"Categorical columns\")\nparser.add_argument(\"--cats_mh\", default=None, help=\"Categorical multihot columns\")\nparser.add_argument(\"--conts\", default=None, help=\"Continuous columns\")\nparser.add_argument(\"--labels\", default=None, help=\"Label columns\")\nparser.add_argument(\"--epochs\", default=1, help=\"Training epochs\")\nargs = parser.parse_args()\n\nhvd.init()\n\ngpu_to_use = hvd.local_rank()\n\nif torch.cuda.is_available():\n    torch.cuda.set_device(gpu_to_use)\n\n\nBASE_DIR = os.path.expanduser(args.dir_in or \"./data/\")\nBATCH_SIZE = args.batch_size or 16384  # Batch Size\nCATEGORICAL_COLUMNS = args.cats or [\"movieId\", \"userId\"]  # Single-hot\nCATEGORICAL_MH_COLUMNS = args.cats_mh or [\"genres\"]  # Multi-hot\nNUMERIC_COLUMNS = args.conts or []\n\n# Output from ETL-with-NVTabular\nTRAIN_PATHS = sorted(glob.glob(os.path.join(BASE_DIR, \"train\", \"*.parquet\")))\n\nproc = nvt.Workflow.load(os.path.join(BASE_DIR, \"workflow/\"))\n\nEMBEDDING_TABLE_SHAPES = nvt.ops.get_embedding_sizes(proc)\n\n\n# TensorItrDataset returns a single batch of x_cat, x_cont, y.\ndef collate_fn(x):\n    return x\n\n\n# Seed with system randomness (or a static seed)\ncupy.random.seed(None)\n\n\ndef seed_fn():\n    \"\"\"\n    Generate consistent dataloader shuffle seeds across workers\n\n    Reseeds each worker's dataloader each epoch to get fresh a shuffle\n    that's consistent across workers.\n    \"\"\"\n\n    max_rand = torch.iinfo(torch.int).max // hvd.size()\n\n    # Generate a seed fragment\n    seed_fragment = cupy.random.randint(0, max_rand)\n\n    # Aggregate seed fragments from all Horovod workers\n    seed_tensor = torch.tensor(seed_fragment)\n    reduced_seed = hvd.allreduce(seed_tensor, name=\"shuffle_seed\", op=hvd.mpi_ops.Sum) % max_rand\n\n    return reduced_seed\n\n\ntrain_dataset = TorchAsyncItr(\n    nvt.Dataset(TRAIN_PATHS),\n    batch_size=BATCH_SIZE,\n    cats=CATEGORICAL_COLUMNS + CATEGORICAL_MH_COLUMNS,\n    conts=NUMERIC_COLUMNS,\n    labels=[\"rating\"],\n    devices=[gpu_to_use],\n    global_size=hvd.size(),\n    global_rank=hvd.rank(),\n    shuffle=True,\n    seed_fn=seed_fn,\n)\ntrain_loader = DLDataLoader(\n    train_dataset, batch_size=None, collate_fn=collate_fn, pin_memory=False, num_workers=0\n)\n\n\nEMBEDDING_TABLE_SHAPES_TUPLE = (\n    {\n        CATEGORICAL_COLUMNS[0]: EMBEDDING_TABLE_SHAPES[CATEGORICAL_COLUMNS[0]],\n        CATEGORICAL_COLUMNS[1]: EMBEDDING_TABLE_SHAPES[CATEGORICAL_COLUMNS[1]],\n    },\n    {CATEGORICAL_MH_COLUMNS[0]: EMBEDDING_TABLE_SHAPES[CATEGORICAL_MH_COLUMNS[0]]},\n)\n\nmodel = Model(\n    embedding_table_shapes=EMBEDDING_TABLE_SHAPES_TUPLE,\n    num_continuous=0,\n    emb_dropout=0.0,\n    layer_hidden_dims=[128, 128, 128],\n    layer_dropout_rates=[0.0, 0.0, 0.0],\n).cuda()\n\nlr_scaler = hvd.size()\n\noptimizer = torch.optim.Adam(model.parameters(), lr=0.01 * lr_scaler)\n\nhvd.broadcast_parameters(model.state_dict(), root_rank=0)\nhvd.broadcast_optimizer_state(optimizer, root_rank=0)\n\noptimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())\n\nfor epoch in range(args.epochs):\n    start = time()\n    print(f\"Training epoch {epoch}\")\n    train_loss, y_pred, y = process_epoch(train_loader, model, train=True, optimizer=optimizer)\n    hvd.join(gpu_to_use)\n    hvd.broadcast_parameters(model.state_dict(), root_rank=0)\n    print(f\"Epoch {epoch:02d}. Train loss: {train_loss:.4f}.\")\n    hvd.join(gpu_to_use)\n    t_final = time() - start\n    total_rows = train_dataset.num_rows_processed\n    print(\n        f\"run_time: {t_final} - rows: {total_rows} - \"\n        f\"epochs: {epoch} - dl_thru: {total_rows / t_final}\"\n    )\n\n\nhvd.join(gpu_to_use)\nif hvd.local_rank() == 0:\n    print(\"Training complete\")\n", "path": "examples/horovod/torch-nvt-horovod.py"}], "after_files": [{"content": "# External dependencies\nimport argparse\nimport glob\nimport os\n\nimport cupy\n\n# we can control how much memory to give tensorflow with this environment variable\n# IMPORTANT: make sure you do this before you initialize TF's runtime, otherwise\n# TF will have claimed all free GPU memory\nos.environ[\"TF_MEMORY_ALLOCATION\"] = \"0.3\"  # fraction of free memory\nimport horovod.tensorflow as hvd  # noqa: E402\nimport tensorflow as tf  # noqa: E402\n\nimport nvtabular as nvt  # noqa: E402\nfrom nvtabular.framework_utils.tensorflow import layers  # noqa: E402\nfrom nvtabular.loader.tensorflow import KerasSequenceLoader  # noqa: E402\n\nparser = argparse.ArgumentParser(description=\"Process some integers.\")\nparser.add_argument(\"--dir_in\", default=None, help=\"Input directory\")\nparser.add_argument(\"--batch_size\", default=None, help=\"batch size\")\nparser.add_argument(\"--cats\", default=None, help=\"categorical columns\")\nparser.add_argument(\"--cats_mh\", default=None, help=\"categorical multihot columns\")\nparser.add_argument(\"--conts\", default=None, help=\"continuous columns\")\nparser.add_argument(\"--labels\", default=None, help=\"continuous columns\")\nargs = parser.parse_args()\n\n\nBASE_DIR = args.dir_in or \"./data/\"\nBATCH_SIZE = int(args.batch_size) or 16384  # Batch Size\nCATEGORICAL_COLUMNS = args.cats or [\"movieId\", \"userId\"]  # Single-hot\nCATEGORICAL_MH_COLUMNS = args.cats_mh or [\"genres\"]  # Multi-hot\nNUMERIC_COLUMNS = args.conts or []\nTRAIN_PATHS = sorted(\n    glob.glob(os.path.join(BASE_DIR, \"train/*.parquet\"))\n)  # Output from ETL-with-NVTabular\nhvd.init()\n\n# Seed with system randomness (or a static seed)\ncupy.random.seed(None)\n\n\ndef seed_fn():\n    \"\"\"\n    Generate consistent dataloader shuffle seeds across workers\n\n    Reseeds each worker's dataloader each epoch to get fresh a shuffle\n    that's consistent across workers.\n    \"\"\"\n    min_int, max_int = tf.int32.limits\n    max_rand = max_int // hvd.size()\n\n    # Generate a seed fragment\n    seed_fragment = cupy.random.randint(0, max_rand).get()\n\n    # Aggregate seed fragments from all Horovod workers\n    seed_tensor = tf.constant(seed_fragment)\n    reduced_seed = hvd.allreduce(seed_tensor, name=\"shuffle_seed\", op=hvd.mpi_ops.Sum) % max_rand\n\n    return reduced_seed\n\n\nproc = nvt.Workflow.load(os.path.join(BASE_DIR, \"workflow/\"))\nEMBEDDING_TABLE_SHAPES = nvt.ops.get_embedding_sizes(proc)\n\ntrain_dataset_tf = KerasSequenceLoader(\n    TRAIN_PATHS,  # you could also use a glob pattern\n    batch_size=BATCH_SIZE,\n    label_names=[\"rating\"],\n    cat_names=CATEGORICAL_COLUMNS + CATEGORICAL_MH_COLUMNS,\n    cont_names=NUMERIC_COLUMNS,\n    engine=\"parquet\",\n    shuffle=True,\n    seed_fn=seed_fn,\n    buffer_size=0.06,  # how many batches to load at once\n    parts_per_chunk=1,\n    global_size=hvd.size(),\n    global_rank=hvd.rank(),\n)\ninputs = {}  # tf.keras.Input placeholders for each feature to be used\nemb_layers = []  # output of all embedding layers, which will be concatenated\nfor col in CATEGORICAL_COLUMNS:\n    inputs[col] = tf.keras.Input(name=col, dtype=tf.int32, shape=(1,))\n# Note that we need two input tensors for multi-hot categorical features\nfor col in CATEGORICAL_MH_COLUMNS:\n    inputs[col + \"__values\"] = tf.keras.Input(name=f\"{col}__values\", dtype=tf.int64, shape=(1,))\n    inputs[col + \"__nnzs\"] = tf.keras.Input(name=f\"{col}__nnzs\", dtype=tf.int64, shape=(1,))\nfor col in CATEGORICAL_COLUMNS + CATEGORICAL_MH_COLUMNS:\n    emb_layers.append(\n        tf.feature_column.embedding_column(\n            tf.feature_column.categorical_column_with_identity(\n                col, EMBEDDING_TABLE_SHAPES[col][0]  # Input dimension (vocab size)\n            ),\n            EMBEDDING_TABLE_SHAPES[col][1],  # Embedding output dimension\n        )\n    )\nemb_layer = layers.DenseFeatures(emb_layers)\nx_emb_output = emb_layer(inputs)\nx = tf.keras.layers.Dense(128, activation=\"relu\")(x_emb_output)\nx = tf.keras.layers.Dense(128, activation=\"relu\")(x)\nx = tf.keras.layers.Dense(128, activation=\"relu\")(x)\nx = tf.keras.layers.Dense(1, activation=\"sigmoid\")(x)\nmodel = tf.keras.Model(inputs=inputs, outputs=x)\nloss = tf.losses.BinaryCrossentropy()\nopt = tf.keras.optimizers.SGD(0.01 * hvd.size())\nopt = hvd.DistributedOptimizer(opt)\ncheckpoint_dir = \"./checkpoints\"\ncheckpoint = tf.train.Checkpoint(model=model, optimizer=opt)\n\n\n@tf.function(experimental_relax_shapes=True)\ndef training_step(examples, labels, first_batch):\n    with tf.GradientTape() as tape:\n        probs = model(examples, training=True)\n        loss_value = loss(labels, probs)\n    # Horovod: add Horovod Distributed GradientTape.\n    tape = hvd.DistributedGradientTape(tape, sparse_as_dense=True)\n    grads = tape.gradient(loss_value, model.trainable_variables)\n    opt.apply_gradients(zip(grads, model.trainable_variables))\n    # Horovod: broadcast initial variable states from rank 0 to all other processes.\n    # This is necessary to ensure consistent initialization of all workers when\n    # training is started with random weights or restored from a checkpoint.\n    #\n    # Note: broadcast should be done after the first gradient step to ensure optimizer\n    # initialization.\n    if first_batch:\n        hvd.broadcast_variables(model.variables, root_rank=0)\n        hvd.broadcast_variables(opt.variables(), root_rank=0)\n    return loss_value\n\n\n# Horovod: adjust number of steps based on number of GPUs.\nfor batch, (examples, labels) in enumerate(train_dataset_tf):\n    loss_value = training_step(examples, labels, batch == 0)\n    if batch % 10 == 0 and hvd.local_rank() == 0:\n        print(\"Step #%d\\tLoss: %.6f\" % (batch, loss_value))\nhvd.join()\n# Horovod: save checkpoints only on worker 0 to prevent other workers from\n# corrupting it.\nif hvd.rank() == 0:\n    checkpoint.save(checkpoint_dir)\n", "path": "examples/horovod/tf_hvd_simple.py"}, {"content": "import argparse\nimport glob\nimport os\nfrom time import time\n\nimport cupy\nimport torch\n\nimport nvtabular as nvt\nfrom nvtabular.framework_utils.torch.models import Model\nfrom nvtabular.framework_utils.torch.utils import process_epoch\nfrom nvtabular.loader.torch import DLDataLoader, TorchAsyncItr\n\n# Horovod must be the last import to avoid conflicts\nimport horovod.torch as hvd  # noqa: E402, isort:skip\n\n\nparser = argparse.ArgumentParser(description=\"Train a multi-gpu model with Torch and Horovod\")\nparser.add_argument(\"--dir_in\", default=None, help=\"Input directory\")\nparser.add_argument(\"--batch_size\", default=None, help=\"Batch size\")\nparser.add_argument(\"--cats\", default=None, help=\"Categorical columns\")\nparser.add_argument(\"--cats_mh\", default=None, help=\"Categorical multihot columns\")\nparser.add_argument(\"--conts\", default=None, help=\"Continuous columns\")\nparser.add_argument(\"--labels\", default=None, help=\"Label columns\")\nparser.add_argument(\"--epochs\", default=1, help=\"Training epochs\")\nargs = parser.parse_args()\n\nhvd.init()\n\ngpu_to_use = hvd.local_rank()\n\nif torch.cuda.is_available():\n    torch.cuda.set_device(gpu_to_use)\n\n\nBASE_DIR = os.path.expanduser(args.dir_in or \"./data/\")\nBATCH_SIZE = int(args.batch_size) or 16384  # Batch Size\nCATEGORICAL_COLUMNS = args.cats or [\"movieId\", \"userId\"]  # Single-hot\nCATEGORICAL_MH_COLUMNS = args.cats_mh or [\"genres\"]  # Multi-hot\nNUMERIC_COLUMNS = args.conts or []\n\n# Output from ETL-with-NVTabular\nTRAIN_PATHS = sorted(glob.glob(os.path.join(BASE_DIR, \"train\", \"*.parquet\")))\n\nproc = nvt.Workflow.load(os.path.join(BASE_DIR, \"workflow/\"))\n\nEMBEDDING_TABLE_SHAPES = nvt.ops.get_embedding_sizes(proc)\n\n\n# TensorItrDataset returns a single batch of x_cat, x_cont, y.\ndef collate_fn(x):\n    return x\n\n\n# Seed with system randomness (or a static seed)\ncupy.random.seed(None)\n\n\ndef seed_fn():\n    \"\"\"\n    Generate consistent dataloader shuffle seeds across workers\n\n    Reseeds each worker's dataloader each epoch to get fresh a shuffle\n    that's consistent across workers.\n    \"\"\"\n\n    max_rand = torch.iinfo(torch.int).max // hvd.size()\n\n    # Generate a seed fragment\n    seed_fragment = cupy.random.randint(0, max_rand)\n\n    # Aggregate seed fragments from all Horovod workers\n    seed_tensor = torch.tensor(seed_fragment)\n    reduced_seed = hvd.allreduce(seed_tensor, name=\"shuffle_seed\", op=hvd.mpi_ops.Sum) % max_rand\n\n    return reduced_seed\n\n\ntrain_dataset = TorchAsyncItr(\n    nvt.Dataset(TRAIN_PATHS),\n    batch_size=BATCH_SIZE,\n    cats=CATEGORICAL_COLUMNS + CATEGORICAL_MH_COLUMNS,\n    conts=NUMERIC_COLUMNS,\n    labels=[\"rating\"],\n    devices=[gpu_to_use],\n    global_size=hvd.size(),\n    global_rank=hvd.rank(),\n    shuffle=True,\n    seed_fn=seed_fn,\n)\ntrain_loader = DLDataLoader(\n    train_dataset, batch_size=None, collate_fn=collate_fn, pin_memory=False, num_workers=0\n)\n\n\nEMBEDDING_TABLE_SHAPES_TUPLE = (\n    {\n        CATEGORICAL_COLUMNS[0]: EMBEDDING_TABLE_SHAPES[CATEGORICAL_COLUMNS[0]],\n        CATEGORICAL_COLUMNS[1]: EMBEDDING_TABLE_SHAPES[CATEGORICAL_COLUMNS[1]],\n    },\n    {CATEGORICAL_MH_COLUMNS[0]: EMBEDDING_TABLE_SHAPES[CATEGORICAL_MH_COLUMNS[0]]},\n)\n\nmodel = Model(\n    embedding_table_shapes=EMBEDDING_TABLE_SHAPES_TUPLE,\n    num_continuous=0,\n    emb_dropout=0.0,\n    layer_hidden_dims=[128, 128, 128],\n    layer_dropout_rates=[0.0, 0.0, 0.0],\n).cuda()\n\nlr_scaler = hvd.size()\n\noptimizer = torch.optim.Adam(model.parameters(), lr=0.01 * lr_scaler)\n\nhvd.broadcast_parameters(model.state_dict(), root_rank=0)\nhvd.broadcast_optimizer_state(optimizer, root_rank=0)\n\noptimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())\n\nfor epoch in range(args.epochs):\n    start = time()\n    print(f\"Training epoch {epoch}\")\n    train_loss, y_pred, y = process_epoch(train_loader, model, train=True, optimizer=optimizer)\n    hvd.join(gpu_to_use)\n    hvd.broadcast_parameters(model.state_dict(), root_rank=0)\n    print(f\"Epoch {epoch:02d}. Train loss: {train_loss:.4f}.\")\n    hvd.join(gpu_to_use)\n    t_final = time() - start\n    total_rows = train_dataset.num_rows_processed\n    print(\n        f\"run_time: {t_final} - rows: {total_rows} - \"\n        f\"epochs: {epoch} - dl_thru: {total_rows / t_final}\"\n    )\n\n\nhvd.join(gpu_to_use)\nif hvd.local_rank() == 0:\n    print(\"Training complete\")\n", "path": "examples/horovod/torch-nvt-horovod.py"}]}