#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import argparse
import itertools
import os
import sys
from dataclasses import dataclass, field
from enum import Enum
from pprint import pprint
from typing import List, Optional

import mlperf_logging.mllog as mllog
import mlperf_logging.mllog.constants as mllog_constants
import torch
import torchmetrics as metrics
from pyre_extensions import none_throws
from torch import distributed as dist
from torch.utils.data import DataLoader
from torchrec import EmbeddingBagCollection
from torchrec.datasets.criteo import DEFAULT_CAT_NAMES, DEFAULT_INT_NAMES
from torchrec.distributed import TrainPipelineSparseDist
from torchrec.distributed.comm import get_local_size
from torchrec.distributed.model_parallel import (
    DistributedModelParallel,
    get_default_sharders,
)
from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
from torchrec.distributed.planner.storage_reservations import HeuristicalStorageReservation
from torchrec.distributed.planner.stats import EmbeddingStats
from torchrec.models.dlrm import DLRM, DLRM_DCN, DLRM_Projection, DLRMTrain
from torchrec.modules.embedding_configs import EmbeddingBagConfig
from torchrec.optim.apply_optimizer_in_backward import apply_optimizer_in_backward
from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizerWrapper
from torchrec.optim.optimizers import in_backward_optimizer_filter
from tqdm import tqdm
import random
import numpy as np
import torch.distributed as dist

# Custom imports for Logging Pipeline
from typing import (
   Callable,
    cast,
    Iterator,
    List,
    Optional,
    Tuple,
    Type,
)
from torchrec.distributed.train_pipeline.utils import (
    In,
    Out,
    TrainPipelineContext,
    _wait_for_batch
)
from torch.autograd.profiler import record_function

# OSS import
try:
    # pyre-ignore[21]
    # @manual=//ai_codesign/benchmarks/dlrm/torchrec_dlrm/data:dlrm_dataloader
    from data.dlrm_dataloader import get_dataloader

    # pyre-ignore[21]
    # @manual=//ai_codesign/benchmarks/dlrm/torchrec_dlrm:lr_scheduler
    from lr_scheduler import LRPolicyScheduler

    # pyre-ignore[21]
    # @manual=//ai_codesign/benchmarks/dlrm/torchrec_dlrm:mlperf_logging_utils
    from mlperf_logging_utils import submission_info

    # pyre-ignore[21]
    # @manual=//ai_codesign/benchmarks/dlrm/torchrec_dlrm:multi_hot
    from multi_hot import Multihot, RestartableMap

    from logging_utils import (
    extract_mlp_weights_and_gradients,
    print_mlp_summary,
    save_mlp_data,
    )

except ImportError:
    pass

# internal import
try:
    from .data.dlrm_dataloader import get_dataloader  # noqa F811
    from .lr_scheduler import LRPolicyScheduler  # noqa F811
    from .mlperf_logging_utils import submission_info  # noqa F811
    from .multi_hot import Multihot, RestartableMap  # noqa F811
except ImportError:
    pass

TRAIN_PIPELINE_STAGES = 3  # Number of stages in TrainPipelineSparseDist.

# Optimizer parameters
ADAGRAD_LR_DECAY = 0
ADAGRAD_INIT_ACC = 0
ADAGRAD_EPS = 1e-8
WEIGHT_DECAY = 0

mllogger = mllog.get_mllogger()

def all_seed(seed: int = 42, deterministic: bool = True):
    """
    Fix all random seeds for reproducibility in distributed training.
    """
    # Python random
    random.seed(seed)
    # NumPy random
    np.random.seed(seed)
    # PyTorch CPU random
    torch.manual_seed(seed)
    # PyTorch GPU random (single- and multi-GPU)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    # cuDNN
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    else:
        torch.backends.cudnn.deterministic = False
        torch.backends.cudnn.benchmark = True

    os.environ["PYTHONHASHSEED"] = str(seed)

    if dist.is_available() and dist.is_initialized():
        rank = dist.get_rank()
        torch.manual_seed(seed + rank)
        torch.cuda.manual_seed(seed + rank)
        np.random.seed(seed + rank)
        random.seed(seed + rank)


# Custom Pipeline Class with logging
class TrainPipelineSparseDist_Log(TrainPipelineSparseDist):
    """
    Custom TrainPipelineSparseDist with logging for each batch processed.
    """

    def __init__(
        self,
        model: torch.nn.Module,
        optimizer: torch.optim.Optimizer,
        device: torch.device,
        execute_all_batches: bool = True,
        apply_jit: bool = False,
        context_type: Type[TrainPipelineContext] = TrainPipelineContext,
        # keep for backward compatibility
        pipeline_postproc: bool = False,
        custom_model_fwd: Optional[
            Callable[[Optional[In]], Tuple[torch.Tensor, Out]]
        ] = None,
    ) -> None:
        super().__init__(model, optimizer, device, execute_all_batches, apply_jit, context_type, pipeline_postproc, custom_model_fwd)

    def progress(self, dataloader_iter: Iterator[In]) -> Out:
        if not self._model_attached:
            self.attach(self._model)

        self.fill_pipeline(dataloader_iter)
        if not self.batches:
            raise StopIteration

        # TODO: Remove once Bulk Eval migrated (needed for bwd compat, this class only)
        self._set_module_context(self.contexts[0])

        if self._model.training:
            with record_function("## zero_grad ##"):
                self._optimizer.zero_grad()

        with record_function("## wait_for_batch ##"):
            _wait_for_batch(cast(In, self.batches[0]), self._data_dist_stream)

        if len(self.batches) >= 2:
            self.start_sparse_data_dist(self.batches[1], self.contexts[1])

        # batch i+2
        self.enqueue_batch(dataloader_iter)

        # forward
        with record_function("## forward ##"):
            losses, output = self._model_fwd(self.batches[0])
        
        if dist.get_rank() == 0:
            print(f" Mean loss: {torch.sum(losses, dim=0):.8f}")

        if len(self.batches) >= 2:
            self.wait_sparse_data_dist(self.contexts[1])

        if self._model.training:
            # backward
            with record_function("## backward ##"):
                torch.sum(losses, dim=0).backward()

            # update
            with record_function("## optimizer ##"):
                self._optimizer.step()

        self.dequeue_batch()
        return output


class InteractionType(Enum):
    ORIGINAL = "original"
    DCN = "dcn"
    PROJECTION = "projection"

    def __str__(self):
        return self.value


def parse_args(argv: List[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="torchrec dlrm example trainer")
    parser.add_argument(
        "--epochs",
        type=int,
        default=1,
        help="number of epochs to train",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=32,
        help="batch size to use for training",
    )
    parser.add_argument(
        "--drop_last_training_batch",
        dest="drop_last_training_batch",
        action="store_true",
        help="Drop the last non-full training batch",
    )
    parser.add_argument(
        "--test_batch_size",
        type=int,
        default=None,
        help="batch size to use for validation and testing",
    )
    parser.add_argument(
        "--limit_train_batches",
        type=int,
        default=None,
        help="number of train batches",
    )
    parser.add_argument(
        "--limit_val_batches",
        type=int,
        default=None,
        help="number of validation batches",
    )
    parser.add_argument(
        "--limit_test_batches",
        type=int,
        default=None,
        help="number of test batches",
    )
    parser.add_argument(
        "--dataset_name",
        type=str,
        default="criteo_1t",
        help="dataset for experiment, current support criteo_1tb, criteo_kaggle",
    )
    parser.add_argument(
        "--num_embeddings",
        type=int,
        default=100_000,
        help="max_ind_size. The number of embeddings in each embedding table. Defaults"
        " to 100_000 if num_embeddings_per_feature is not supplied.",
    )
    parser.add_argument(
        "--num_embeddings_per_feature",
        type=str,
        default=None,
        help="Comma separated max_ind_size per sparse feature. The number of embeddings"
        " in each embedding table. 26 values are expected for the Criteo dataset.",
    )
    parser.add_argument(
        "--dense_arch_layer_sizes",
        type=str,
        default="512,256,64",
        help="Comma separated layer sizes for dense arch.",
    )
    parser.add_argument(
        "--over_arch_layer_sizes",
        type=str,
        default="512,512,256,1",
        help="Comma separated layer sizes for over arch.",
    )
    parser.add_argument(
        "--embedding_dim",
        type=int,
        default=64,
        help="Size of each embedding.",
    )
    parser.add_argument(
        "--interaction_branch1_layer_sizes",
        type=str,
        default="2048,2048",
        help="Comma separated layer sizes for interaction branch1 (only on dlrm with projection).",
    )
    parser.add_argument(
        "--interaction_branch2_layer_sizes",
        type=str,
        default="2048,2048",
        help="Comma separated layer sizes for interaction branch2 (only on dlrm with projection).",
    )
    parser.add_argument(
        "--dcn_num_layers",
        type=int,
        default=3,
        help="Number of DCN layers in interaction layer (only on dlrm with DCN).",
    )
    parser.add_argument(
        "--dcn_low_rank_dim",
        type=int,
        default=512,
        help="Low rank dimension for DCN in interaction layer (only on dlrm with DCN).",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed for reproducibility.",
    )
    parser.add_argument(
        "--pin_memory",
        dest="pin_memory",
        action="store_true",
        help="Use pinned memory when loading data.",
    )
    parser.add_argument(
        "--mmap_mode",
        dest="mmap_mode",
        action="store_true",
        help="--mmap_mode mmaps the dataset."
        " That is, the dataset is kept on disk but is accessed as if it were in memory."
        " --mmap_mode is intended mostly for faster debugging. Use --mmap_mode to bypass"
        " preloading the dataset when preloading takes too long or when there is "
        " insufficient memory available to load the full dataset.",
    )
    parser.add_argument(
        "--in_memory_binary_criteo_path",
        type=str,
        default=None,
        help="Directory path containing the Criteo dataset npy files.",
    )
    parser.add_argument(
        "--synthetic_multi_hot_criteo_path",
        type=str,
        default=None,
        help="Directory path containing the MLPerf v2 synthetic multi-hot dataset npz files.",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=15.0,
        help="Learning rate.",
    )
    parser.add_argument(
        "--shuffle_batches",
        dest="shuffle_batches",
        action="store_true",
        help="Shuffle each batch during training.",
    )
    parser.add_argument(
        "--shuffle_training_set",
        dest="shuffle_training_set",
        action="store_true",
        help="Shuffle the training set in memory. This will override mmap_mode",
    )
    parser.add_argument(
        "--validation_freq_within_epoch",
        type=int,
        default=None,
        help="Frequency at which validation will be run within an epoch.",
    )
    parser.add_argument(
        "--validation_auroc",
        type=float,
        default=None,
        help="Validation AUROC threshold to stop training once reached.",
    )
    parser.add_argument(
        "--evaluate_on_epoch_end",
        action="store_true",
        help="Evaluate using validation set on each epoch end.",
    )
    parser.add_argument(
        "--evaluate_on_training_end",
        action="store_true",
        help="Evaluate using test set on training end.",
    )
    parser.add_argument(
        "--adagrad",
        dest="adagrad",
        action="store_true",
        help="Flag to determine if adagrad optimizer should be used.",
    )
    parser.add_argument(
        "--interaction_type",
        type=InteractionType,
        choices=list(InteractionType),
        default=InteractionType.ORIGINAL,
        help="Determine the interaction type to be used (original, dcn, or projection)"
        " default is original DLRM with pairwise dot product",
    )
    parser.add_argument(
        "--collect_multi_hot_freqs_stats",
        dest="collect_multi_hot_freqs_stats",
        action="store_true",
        help="Flag to determine whether to collect stats on freq of embedding access.",
    )
    parser.add_argument(
        "--multi_hot_sizes",
        type=str,
        default=None,
        help="Comma separated multihot size per sparse feature. 26 values are expected for the Criteo dataset.",
    )
    parser.add_argument(
        "--multi_hot_distribution_type",
        type=str,
        choices=["uniform", "pareto"],
        default=None,
        help="Multi-hot distribution options.",
    )
    parser.add_argument(
        "--lr_warmup_steps",
        type=int,
        default=0,
    )
    parser.add_argument(
        "--lr_decay_start",
        type=int,
        default=0,
    )
    parser.add_argument(
        "--lr_decay_steps",
        type=int,
        default=0,
    )
    parser.add_argument(
        "--print_lr",
        action="store_true",
        help="Print learning rate every iteration.",
    )
    parser.add_argument(
        "--allow_tf32",
        action="store_true",
        help="Enable TensorFloat-32 mode for matrix multiplications on A100 (or newer) GPUs.",
    )
    parser.add_argument(
        "--print_sharding_plan",
        action="store_true",
        help="Print the sharding plan used for each embedding table.",
    )
    parser.add_argument(
        "--print_progress",
        action="store_true",
        help="Print tqdm progress bar during training and evaluation.",
    )
    parser.add_argument(
        "--log_params",
        action="store_true",
        help="Save parameters and gradients of the MLP layers for step 1.",
    )
    parser.add_argument(
        "--profile_dir",
        type=str,
        default=None,
        help="Directory for savi torch traces.",
    )
    return parser.parse_args(argv)


def _evaluate(
    limit_batches: Optional[int],
    eval_pipeline: TrainPipelineSparseDist,
    eval_dataloader: DataLoader,
    stage: str,
    epoch_num: float,
    log_eval_samples: bool,
    print_progress: bool,
) -> float:
    """
    Evaluates model. Computes and prints AUROC. Helper function for train_val_test.

    Args:
        limit_batches (Optional[int]): Limits the dataloader to the first `limit_batches` batches.
        eval_pipeline (TrainPipelineSparseDist): pipelined model.
        eval_dataloader (DataLoader): Dataloader for either the validation set or test set.
        stage (str): "val" or "test".
        epoch_num (float): Iterations passed as epoch fraction (for logging purposes).
        log_eval_samples (bool): Whether to print mllog with the number of samples.
        print_progress (bool): Whether to print tqdm progress bar.

    Returns:
        float: auroc result
    """
    is_rank_zero = dist.get_rank() == 0
    if is_rank_zero:
        pbar = tqdm(
            iter(int, 1),
            desc=f"Evaluating {stage} set",
            total=len(eval_dataloader),
            disable=not print_progress,
        )
        mllogger.start(
            key=mllog_constants.EVAL_START,
            metadata={mllog_constants.EPOCH_NUM: epoch_num},
        )

    eval_pipeline._model.eval()
    device = eval_pipeline._device

    # Set eval_pipeline._connected to False to cause the pipeline to refill with new batches as if it were newly created and empty.
    eval_pipeline._connected = False

    iterator = itertools.islice(iter(eval_dataloader), limit_batches)
    # Two filler batches are appended to the end of the iterator to keep the pipeline active while the
    # last two remaining batches are still in progress awaiting results.
    two_filler_batches = itertools.islice(
        iter(eval_dataloader), TRAIN_PIPELINE_STAGES - 1
    )
    iterator = itertools.chain(iterator, two_filler_batches)

    auroc = metrics.AUROC(task='binary').to(device)

    with torch.no_grad():
        while True:
            try:
                _loss, logits, labels = eval_pipeline.progress(iterator)
                preds = torch.sigmoid(logits)
                auroc(preds, labels)
                if is_rank_zero:
                    pbar.update(1)
            except StopIteration:
                # Dataset traversal complete
                break

    auroc_result = auroc.compute().item()
    num_samples = torch.tensor(sum(map(len, auroc.target)), device=device)
    dist.reduce(num_samples, 0, op=dist.ReduceOp.SUM)
    num_samples = num_samples.item()

    if is_rank_zero:
        print(f"AUROC over {stage} set: {auroc_result}.")
        print(f"Number of {stage} samples: {num_samples}")
        mllogger.event(
            key=mllog_constants.EVAL_ACCURACY,
            value=auroc_result,
            metadata={mllog_constants.EPOCH_NUM: epoch_num},
        )
        mllogger.end(
            key=mllog_constants.EVAL_STOP,
            metadata={mllog_constants.EPOCH_NUM: epoch_num},
        )
        if log_eval_samples:
            mllogger.event(
                key=mllog_constants.EVAL_SAMPLES,
                value=num_samples,
                metadata={mllog_constants.EPOCH_NUM: epoch_num},
            )
    return auroc_result


def _train(
    train_pipeline: TrainPipelineSparseDist_Log,
    val_pipeline: TrainPipelineSparseDist,
    train_dataloader: DataLoader,
    val_dataloader: DataLoader,
    epoch: int,
    lr_scheduler,
    print_lr: bool,
    validation_freq: Optional[int],
    validation_auroc: Optional[float],
    limit_train_batches: Optional[int],
    limit_val_batches: Optional[int],
    print_progress: bool,
    log_params: bool = False,
) -> bool:
    """
    Trains model for 1 epoch. Helper function for train_val_test.

    Args:
        train_pipeline (TrainPipelineSparseDist): pipelined model used for training.
        val_pipeline (TrainPipelineSparseDist): pipelined model used for validation.
        train_dataloader (DataLoader): Training set's dataloader.
        val_dataloader (DataLoader): Validation set's dataloader.
        epoch (int): The number of complete passes through the training set so far.
        lr_scheduler (LRPolicyScheduler): Learning rate scheduler.
        print_lr (bool): Whether to print the learning rate every training step.
        validation_freq (Optional[int]): The number of training steps between validation runs within an epoch.
        validation_auroc (Optional[float]): AUROC level desired for stopping training.
        limit_train_batches (Optional[int]): Limits the training set to the first `limit_train_batches` batches.
        limit_val_batches (Optional[int]): Limits the validation set to the first `limit_val_batches` batches.
        print_progress (bool): Whether to print tqdm progress bar.

    Returns:
        bool: Whether the validation_auroc threshold is reached.
    """
    train_pipeline._model.train()

    # Set train_pipeline._connected to False to cause the pipeline to refill with new batches as if it were newly created and empty.
    train_pipeline._connected = False

    iterator = itertools.islice(iter(train_dataloader), limit_train_batches)
    # Two filler batches are appended to the end of the iterator to keep the pipeline active while the
    # last two remaining batches are still in progress awaiting results.
    two_filler_batches = itertools.islice(
        iter(train_dataloader), TRAIN_PIPELINE_STAGES - 1
    )
    iterator = itertools.chain(iterator, two_filler_batches)

    is_rank_zero = dist.get_rank() == 0
    if is_rank_zero:
        pbar = tqdm(
            iter(int, 1),
            desc=f"Epoch {epoch}",
            total=len(train_dataloader),
            disable=not print_progress,
        )

    it = 1
    is_success = False
    is_first_eval = True
    for it in itertools.count(1):
        try:
            if is_rank_zero and print_lr:
                for i, g in enumerate(train_pipeline._optimizer.param_groups):
                    print(f"lr: {it} {i} {g['lr']:.6f}")
            train_pipeline.progress(iterator)
            lr_scheduler.step()
            if is_rank_zero:
                pbar.update(1)
            if validation_freq and it % validation_freq == 0:
                epoch_num = epoch + it / len(train_dataloader)
                auroc_result = _evaluate(
                    limit_val_batches,
                    val_pipeline,
                    val_dataloader,
                    "val",
                    epoch_num,
                    is_first_eval and epoch == 0,
                    print_progress,
                )
                is_first_eval = False
                if validation_auroc is not None and auroc_result >= validation_auroc:
                    dist.barrier()
                    if is_rank_zero:
                        mllogger.end(
                            key=mllog_constants.RUN_STOP,
                            metadata={
                                mllog_constants.STATUS: mllog_constants.SUCCESS,
                                mllog_constants.EPOCH_NUM: epoch_num,
                            },
                        )
                    is_success = True
                    break
                train_pipeline._model.train()

            if log_params and it == 1 and is_rank_zero:
                mlp_data = extract_mlp_weights_and_gradients(train_pipeline)
                print_mlp_summary(mlp_data)      
                save_mlp_data(mlp_data, f'model_param_data.pt')

        except StopIteration:
            # Dataset traversal complete
            break

    if is_rank_zero:
        print("Total number of iterations:", it - 1)

    return is_success


@dataclass
class TrainValTestResults:
    val_aurocs: List[float] = field(default_factory=list)
    test_auroc: Optional[float] = None


def train_val_test(
    args: argparse.Namespace,
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    train_dataloader: DataLoader,
    val_dataloader: DataLoader,
    test_dataloader: DataLoader,
    lr_scheduler: LRPolicyScheduler,
) -> TrainValTestResults:
    """
    Train/validation/test loop.

    Args:
        args (argparse.Namespace): parsed command line args.
        model (torch.nn.Module): model to train.
        optimizer (torch.optim.Optimizer): optimizer to use.
        device (torch.device): device to use.
        train_dataloader (DataLoader): Training set's dataloader.
        val_dataloader (DataLoader): Validation set's dataloader.
        test_dataloader (DataLoader): Test set's dataloader.
        lr_scheduler (LRPolicyScheduler): Learning rate scheduler.

    Returns:
        TrainValTestResults.
    """
    results = TrainValTestResults()
    train_pipeline = TrainPipelineSparseDist_Log(model, optimizer, device)
    val_pipeline = TrainPipelineSparseDist(model, optimizer, device)
    test_pipeline = TrainPipelineSparseDist(model, optimizer, device)

    is_rank_zero = dist.get_rank() == 0

    epoch = 0
    is_success = False
    for epoch in range(args.epochs):
        if is_rank_zero:
            mllogger.start(
                key=mllog_constants.EPOCH_START,
                metadata={mllog_constants.EPOCH_NUM: epoch},
            )
        is_success = _train(
            train_pipeline,
            val_pipeline,
            train_dataloader,
            val_dataloader,
            epoch,
            lr_scheduler,
            args.print_lr,
            args.validation_freq_within_epoch,
            args.validation_auroc,
            args.limit_train_batches,
            args.limit_val_batches,
            args.print_progress,
            args.log_params
        )
        if args.evaluate_on_epoch_end:
            val_auroc = _evaluate(
                args.limit_val_batches,
                val_pipeline,
                val_dataloader,
                "val",
                epoch + 1,
                False,
                args.print_progress,
            )
            results.val_aurocs.append(val_auroc)
        if is_rank_zero:
            mllogger.end(
                key=mllog_constants.EPOCH_STOP,
                metadata={mllog_constants.EPOCH_NUM: epoch},
            )
        if is_success:
            break

    dist.barrier()
    if is_rank_zero and not is_success:
        # Run status "aborted" is reported in the case AUROC threshold is not met
        mllogger.end(
            key=mllog_constants.RUN_STOP,
            metadata={
                mllog_constants.STATUS: mllog_constants.ABORTED,
                mllog_constants.EPOCH_NUM: epoch + 1,
            },
        )

    if args.evaluate_on_training_end:
        test_auroc = _evaluate(
            args.limit_test_batches,
            test_pipeline,
            test_dataloader,
            "test",
            epoch + 1,
            False,
            args.print_progress,
        )
        results.test_auroc = test_auroc

    return results


def main(argv: List[str]) -> None:
    """
    Trains, validates, and tests a Deep Learning Recommendation Model (DLRM)
    (https://arxiv.org/abs/1906.00091). The DLRM model contains both data parallel
    components (e.g. multi-layer perceptrons & interaction arch) and model parallel
    components (e.g. embedding tables). The DLRM model is pipelined so that dataloading,
    data-parallel to model-parallel comms, and forward/backward are overlapped. Can be
    run with either a random dataloader or an in-memory Criteo 1 TB click logs dataset
    (https://ailab.criteo.com/download-criteo-1tb-click-logs-dataset/).

    Args:
        argv (List[str]): command line args.

    Returns:
        None.
    """

    # The reference implementation does not clear the cache currently
    # but the submissions are required to do so
    mllogger.event(key=mllog_constants.CACHE_CLEAR, value=True)
    mllogger.start(key=mllog_constants.INIT_START)

    args = parse_args(argv)
    for name, val in vars(args).items():
        try:
            vars(args)[name] = list(map(int, val.split(",")))
        except (ValueError, AttributeError):
            pass

    torch.backends.cuda.matmul.allow_tf32 = args.allow_tf32

    if args.multi_hot_sizes is not None:
        assert (
            args.num_embeddings_per_feature is not None
            and len(args.multi_hot_sizes) == len(args.num_embeddings_per_feature)
            or args.num_embeddings_per_feature is None
            and len(args.multi_hot_sizes) == len(DEFAULT_CAT_NAMES)
        ), "--multi_hot_sizes must be a comma delimited list the same size as the number of embedding tables."
    assert (
        args.in_memory_binary_criteo_path is None
        or args.synthetic_multi_hot_criteo_path is None
    ), "--in_memory_binary_criteo_path and --synthetic_multi_hot_criteo_path are mutually exclusive CLI arguments."
    assert (
        args.multi_hot_sizes is None or args.synthetic_multi_hot_criteo_path is None
    ), "--multi_hot_sizes is used to convert 1-hot to multi-hot. It's inapplicable with --synthetic_multi_hot_criteo_path."
    assert (
        args.multi_hot_distribution_type is None
        or args.synthetic_multi_hot_criteo_path is None
    ), "--multi_hot_distribution_type is used to convert 1-hot to multi-hot. It's inapplicable with --synthetic_multi_hot_criteo_path."

    rank = int(os.environ["LOCAL_RANK"])
    if torch.cuda.is_available():
        device: torch.device = torch.device(f"cuda:{rank}")
        backend = "nccl"
        torch.cuda.set_device(device)
    else:
        device: torch.device = torch.device("cpu")
        backend = "gloo"
    dist.init_process_group(backend=backend)
    all_seed(args.seed)
    
    is_rank_zero = dist.get_rank() == 0
    if is_rank_zero:
        pprint(vars(args))
        submission_info(mllogger, mllog_constants.DLRM_DCNv2, "reference_implementation")
        mllogger.event(
            key=mllog_constants.GLOBAL_BATCH_SIZE,
            value=dist.get_world_size() * args.batch_size,
        )
        mllogger.event(
            key=mllog_constants.GRADIENT_ACCUMULATION_STEPS,
            value=1,  # Gradient accumulation is not supported in the reference implementation
        )
        mllogger.event(
            key=mllog_constants.SEED,
            value=args.seed,  # Seeding model is not supported in the reference implementation
        )

    if args.num_embeddings_per_feature is not None:
        args.num_embeddings = None

    # Sets default limits for random dataloader iterations when left unspecified.
    if (
        args.in_memory_binary_criteo_path is None
        and args.synthetic_multi_hot_criteo_path is None
    ):
        for split in ["train", "val", "test"]:
            attr = f"limit_{split}_batches"
            if getattr(args, attr) is None:
                setattr(args, attr, 10)

    eb_configs = [
        EmbeddingBagConfig(
            name=f"t_{feature_name}",
            embedding_dim=args.embedding_dim,
            num_embeddings=none_throws(args.num_embeddings_per_feature)[feature_idx]
            if args.num_embeddings is None
            else args.num_embeddings,
            feature_names=[feature_name],
        )
        for feature_idx, feature_name in enumerate(DEFAULT_CAT_NAMES)
    ]
    sharded_module_kwargs = {}
    if args.over_arch_layer_sizes is not None:
        sharded_module_kwargs["over_arch_layer_sizes"] = args.over_arch_layer_sizes

    if args.interaction_type == InteractionType.ORIGINAL:
        dlrm_model = DLRM(
            embedding_bag_collection=EmbeddingBagCollection(
                tables=eb_configs, device=torch.device("meta")
            ),
            dense_in_features=len(DEFAULT_INT_NAMES),
            dense_arch_layer_sizes=args.dense_arch_layer_sizes,
            over_arch_layer_sizes=args.over_arch_layer_sizes,
            dense_device=device,
        )
    elif args.interaction_type == InteractionType.DCN:
        dlrm_model = DLRM_DCN(
            embedding_bag_collection=EmbeddingBagCollection(
                tables=eb_configs, device=torch.device("meta")
            ),
            dense_in_features=len(DEFAULT_INT_NAMES),
            dense_arch_layer_sizes=args.dense_arch_layer_sizes,
            over_arch_layer_sizes=args.over_arch_layer_sizes,
            dcn_num_layers=args.dcn_num_layers,
            dcn_low_rank_dim=args.dcn_low_rank_dim,
            dense_device=device,
        )
    elif args.interaction_type == InteractionType.PROJECTION:
        dlrm_model = DLRM_Projection(
            embedding_bag_collection=EmbeddingBagCollection(
                tables=eb_configs, device=torch.device("meta")
            ),
            dense_in_features=len(DEFAULT_INT_NAMES),
            dense_arch_layer_sizes=args.dense_arch_layer_sizes,
            over_arch_layer_sizes=args.over_arch_layer_sizes,
            interaction_branch1_layer_sizes=args.interaction_branch1_layer_sizes,
            interaction_branch2_layer_sizes=args.interaction_branch2_layer_sizes,
            dense_device=device,
        )
    else:
        raise ValueError(
            "Unknown interaction option set. Should be original, dcn, or projection."
        )

    train_model = DLRMTrain(dlrm_model)
    embedding_optimizer = torch.optim.Adagrad if args.adagrad else torch.optim.SGD
    # This will apply the Adagrad optimizer in the backward pass for the embeddings (sparse_arch). This means that
    # the optimizer update will be applied in the backward pass, in this case through a fused op.
    # TorchRec will use the FBGEMM implementation of EXACT_ADAGRAD. For GPU devices, a fused CUDA kernel is invoked. For CPU, FBGEMM_GPU invokes CPU kernels
    # https://github.com/pytorch/FBGEMM/blob/2cb8b0dff3e67f9a009c4299defbd6b99cc12b8f/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py#L676-L678

    # Note that lr_decay, weight_decay and initial_accumulator_value for Adagrad optimizer in FBGEMM v0.3.2
    # cannot be specified below. This equivalently means that all these parameters are hardcoded to zero.
    optimizer_kwargs = {"lr": args.learning_rate}
    if args.adagrad:
        optimizer_kwargs["eps"] = ADAGRAD_EPS

    apply_optimizer_in_backward(
        embedding_optimizer,
        train_model.model.sparse_arch.parameters(),
        optimizer_kwargs,
    )

    # Set device properties
    device_properties = str(torch.cuda.get_device_properties(0))
    if 'gfx942' in device_properties:
        hbm_cap = 192 * 1024 * 1024 * 1024  # 192GB MI300X HBM
        ddr_cap = 1024 * 1024 * 1024 * 1024  # 1TB host memory per GPU (System Specific)
        hbm_mem_bw = 5.3 * 1024 * 1024 * 1024 * 1024 / 1000  # 5.3 TB/s MI300X
        ddr_mem_bw = 0.8 * 460.8 * 1024 * 1024 * 1024 / 1000 # ~370 GB/s (System Specific. egs. 80% of theoretical 460.8 GB/s AMD EPYC™ 9654). Refer to https://www.amd.com/en/products/processors/server/epyc/4th-generation-9004-and-8004-series/ for other AMD EPYC 4th Gen Processors.
        hbm_to_ddr_mem_bw = 128 * 1024 * 1024 * 1024 / 1000  # 128 GB/s (PCIe Gen5 x16)
        intra_host_bw = 0.8 * 7 * 64 * 1024 * 1024 * 1024 / 1000 # ~336 GB/s (80% of 7x64 GB/s using AMD's xGMI)

    elif 'gfx944' in device_properties:
        hbm_cap = 256 * 1024 * 1024 * 1024  # 256GB MI325X HBM
        ddr_cap = 1024 * 1024 * 1024 * 1024  # 1TB host memory per GPU (Node Specific)
        hbm_mem_bw = 6 * 1024 * 1024 * 1024 * 1024 / 1000  # 6 TB/s MI325X
        ddr_mem_bw = 0.8 * 460.8 * 1024 * 1024 * 1024 / 1000 # ~370 GB/s (System Specific. egs. 80% of theoretical 460.8 GB/s AMD EPYC™ 9654). Refer to https://www.amd.com/en/products/processors/server/epyc/4th-generation-9004-and-8004-series/ for other AMD EPYC 4th Gen Processors.
        hbm_to_ddr_mem_bw = 128 * 1024 * 1024 * 1024 / 1000  # 128 GB/s (PCIe Gen5 x16)
        intra_host_bw = 0.8 * 7 * 64 * 1024 * 1024 * 1024 / 1000 # ~336 GB/s (80% of 7x64 GB/s using AMD's xGMI)

    elif 'gfx950' in device_properties:
        hbm_cap = 288 * 1024 * 1024 * 1024  # 288GB MI350X HBM
        ddr_cap = 1024 * 1024 * 1024 * 1024  # 1TB host memory per GPU (Node Specific)
        hbm_mem_bw = 8 * 1024 * 1024 * 1024 * 1024 / 1000  # 8 TB/s MI350X
        ddr_mem_bw = 0.8 * 614 * 1024 * 1024 * 1024 / 1000 # ~491 GB/s (System Specific. egs. 80% of theoretical 614 GB/s AMD EPYC™ 9655). Refer to https://www.amd.com/en/products/processors/server/epyc/9005-series.html for other AMD EPYC 5th Gen Processors.
        hbm_to_ddr_mem_bw = 128 * 1024 * 1024 * 1024 / 1000  # 128 GB/s (PCIe Gen5 x16)
        intra_host_bw = 0.8 * 7 * 76.8 * 1024 * 1024 * 1024 / 1000 # ~427 GB/s (80% of 7x76.8 GB/s using AMD's xGMI)
    
    else:
        return RuntimeError(f"Unknown device {device_properties}")

    topology = Topology(
        local_world_size=get_local_size(),
        world_size=dist.get_world_size(),
        compute_device=device.type,
        hbm_cap=hbm_cap,
        ddr_cap=ddr_cap,
        hbm_mem_bw=hbm_mem_bw,
        ddr_mem_bw=ddr_mem_bw,
        hbm_to_ddr_mem_bw=hbm_to_ddr_mem_bw,
        intra_host_bw=intra_host_bw,
        )
    
    planner = EmbeddingShardingPlanner(
        topology = topology,
        batch_size=args.batch_size,
        storage_reservation=HeuristicalStorageReservation(percentage=0.10),
        stats=EmbeddingStats()
    )
    
    plan = planner.collective_plan(
        train_model, get_default_sharders(), dist.GroupMember.WORLD
    )

    model = DistributedModelParallel(
        module=train_model,
        device=device,
        plan=plan,
    )
    if is_rank_zero and args.print_sharding_plan:
        for collectionkey, plans in model._plan.plan.items():
            print(collectionkey)
            for table_name, plan in plans.items():
                print(table_name, "\n", plan, "\n")

    def optimizer_with_params():
        if args.adagrad:
            return lambda params: torch.optim.Adagrad(
                params,
                lr=args.learning_rate,
                lr_decay=ADAGRAD_LR_DECAY,
                weight_decay=WEIGHT_DECAY,
                initial_accumulator_value=ADAGRAD_INIT_ACC,
                eps=ADAGRAD_EPS,
            )
        else:
            return lambda params: torch.optim.SGD(
                params,
                lr=args.learning_rate,
                weight_decay=WEIGHT_DECAY,
            )

    dense_optimizer = KeyedOptimizerWrapper(
        dict(in_backward_optimizer_filter(model.named_parameters())),
        optimizer_with_params(),
    )
    optimizer = CombinedOptimizer([model.fused_optimizer, dense_optimizer])
    lr_scheduler = LRPolicyScheduler(
        optimizer, args.lr_warmup_steps, args.lr_decay_start, args.lr_decay_steps
    )

    if is_rank_zero:
        mllogger.event(
            key=mllog_constants.OPT_NAME,
            value=mllog_constants.ADAGRAD if args.adagrad else mllog_constants.SGD,
        )
        mllogger.event(
            key=mllog_constants.OPT_BASE_LR,
            value=args.learning_rate,
        )
        mllogger.event(
            key=mllog_constants.OPT_ADAGRAD_LR_DECAY,
            value=ADAGRAD_LR_DECAY,
        )
        mllogger.event(
            key=mllog_constants.OPT_WEIGHT_DECAY,
            value=WEIGHT_DECAY,
        )
        mllogger.event(
            key=mllog_constants.OPT_ADAGRAD_INITIAL_ACCUMULATOR_VALUE,
            value=ADAGRAD_INIT_ACC,
        )
        mllogger.event(
            key=mllog_constants.OPT_ADAGRAD_EPSILON,
            value=ADAGRAD_EPS,
        )
        mllogger.event(
            key=mllog_constants.OPT_LR_WARMUP_STEPS,
            value=args.lr_warmup_steps,
        )
        mllogger.event(
            key=mllog_constants.OPT_LR_DECAY_START_STEP,
            value=args.lr_decay_start,
        )
        mllogger.event(
            key=mllog_constants.OPT_LR_DECAY_STEPS,
            value=args.lr_decay_steps,
        )

    dist.barrier()
    if is_rank_zero:
        mllogger.start(key=mllog_constants.INIT_STOP)
    dist.barrier()
    if is_rank_zero:
        mllogger.start(key=mllog_constants.RUN_START)
    dist.barrier()

    train_dataloader = get_dataloader(args, backend, "train")
    val_dataloader = get_dataloader(args, backend, "val")
    test_dataloader = get_dataloader(args, backend, "test")

    if is_rank_zero:
        mllogger.event(
            key=mllog_constants.TRAIN_SAMPLES,
            value=dist.get_world_size() * len(train_dataloader) * args.batch_size,
        )

    if args.multi_hot_sizes is not None:
        multihot = Multihot(
            args.multi_hot_sizes,
            args.num_embeddings_per_feature,
            args.batch_size,
            collect_freqs_stats=args.collect_multi_hot_freqs_stats,
            dist_type=args.multi_hot_distribution_type,
        )
        multihot.pause_stats_collection_during_val_and_test(model)
        train_dataloader = RestartableMap(
            multihot.convert_to_multi_hot, train_dataloader
        )
        val_dataloader = RestartableMap(multihot.convert_to_multi_hot, val_dataloader)
        test_dataloader = RestartableMap(multihot.convert_to_multi_hot, test_dataloader)
    train_val_test(
        args,
        model,
        optimizer,
        device,
        train_dataloader,
        val_dataloader,
        test_dataloader,
        lr_scheduler,
    )
    if args.collect_multi_hot_freqs_stats:
        multihot.save_freqs_stats()


if __name__ == "__main__":
    main(sys.argv[1:])
