{ "source_file": "0930", "source_platform": "Medium", "author": "Chaim Rand", "title_listed": "Optimizing Data Transfer in Distributed AI/ML Training Workloads", "keywords": "", "domain": "数据处理", "content_type": "深度技术解析", "is_paid_flag": "否", "url": "https://towardsdatascience.com/optimizing-data-transfer-in-distributed-ai-ml-training-workloads/", "publish_date": "2026-01-23", "corpus_stage": "", "notes": "", "id": "f29ca4fd35ae", "collected_at": "2026-10-02T10:59:01", "title": "Optimizing Data Transfer in Distributed AI/ML Training Workloads | Towards Data Science", "body_markdown": "![Revisit consent button](../images/body-df29beacd7496220cc94.svg)\nWe value your privacy\nWe use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking \"Accept All\", you consent to our use of cookies.\nCustomise Reject All Accept All\nCustomise Consent Preferences ![](../images/body-2159bb67f3de8ad870eb.svg)\nWe use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below.\nThe cookies that are categorised as \"Necessary\" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more\nNecessaryAlways Active\nNecessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data.\n * Cookie\nBCTempID\n * Duration\n10 minutes\n * Description\nNo description available.\n\n\n * Cookie\n__cf_bm\n * Duration\n1 hour\n * Description\nThis cookie, set by Cloudflare, is used to support Cloudflare Bot Management. \n\n\n * Cookie\nAWSALBCORS\n * Duration\n7 days\n * Description\nAmazon Web Services set this cookie for load balancing.\n\n\n * Cookie\n_cfuvid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to track users across sessions to optimize user experience by maintaining session consistency and providing personalized services\n\n\n * Cookie\nli_gc\n * Duration\n6 months\n * Description\nLinkedin set this cookie for storing visitor's consent regarding using cookies for non-essential purposes.\n\n\n * Cookie\n__hssrc\n * Duration\nsession\n * Description\nThis cookie is set by Hubspot whenever it changes the session cookie. The __hssrc cookie set to 1 indicates that the user has restarted the browser, and if the cookie does not exist, it is assumed to be a new session.\n\n\n * Cookie\n__hssc\n * Duration\n1 hour\n * Description\nHubSpot sets this cookie to keep track of sessions and to determine if HubSpot should increment the session number and timestamps in the __hstc cookie.\n\n\n * Cookie\nwpEmojiSettingsSupports\n * Duration\nsession\n * Description\nWordPress sets this cookie when a user interacts with emojis on a WordPress site. It helps determine if the user's browser can display emojis properly.\n\n\n * Cookie\nBCSessionID\n * Duration\n1 year 1 month 4 days\n * Description\nBlueconic sets this cookie as a unique identifier for the BlueConic profile.\n\n\n * Cookie\n_octo\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nlogged_in\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\n__Secure-YEC\n * Duration\npast\n * Description\nYouTube sets this cookie to stores the user's video player preferences using embedded YouTube video\n\n\n * Cookie\n__eoi\n * Duration\n6 months\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTGCORS\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\nlogin-status-p\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTG\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\ncsrf_token\n * Duration\nsession\n * Description\nNo description available.\n\n\n * Cookie\ntoken_v2\n * Duration\n1 day\n * Description\nDescription is currently not available.\n\n\n * Cookie\nD\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nPHPSESSID\n * Duration\nsession\n * Description\nThis cookie is native to PHP applications. The cookie stores and identifies a user's unique session ID to manage user sessions on the website. The cookie is a session cookie and will be deleted when all the browser windows are closed.\n\n\n * Cookie\nVISITOR_PRIVACY_METADATA\n * Duration\n6 months\n * Description\nYouTube sets this cookie to store the user's cookie consent state for the current domain. \n\n\n * Cookie\ncookietest\n * Duration\nsession\n * Description\nThe cookietest cookie is typically used to determine whether the user's browser accepts cookies, essential for website functionality and user experience.\n\n\n * Cookie\n__Host-airtable-session\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\n__Host-airtable-session.sig\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\nm\n * Duration\n1 year 1 month 4 days\n * Description\nStripe sets this cookie for fraud prevention purposes. It identifies the device used to access the website, allowing the website to be formatted accordingly.\n\n\n * Cookie\nBIGipServer*\n * Duration\nsession\n * Description\nMarketo sets this cookie to collect information about the user's online activity and build a profile about their interests to provide advertisements relevant to the user. \n\n\n * Cookie\n__cfruid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to identify trusted web traffic.\n\n\n * Cookie\n_GRECAPTCHA\n * Duration\n6 months\n * Description\nGoogle Recaptcha service sets this cookie to identify bots to protect the website against malicious spam attacks.\n\n\n * Cookie\n__Secure-YNID\n * Duration\n6 months\n * Description\nGoogle cookie used to protect user security and prevent fraud, especially during the login process.\n\n\n * Cookie\ncookieyes-consent\n * Duration\n1 year\n * Description\nCookieYes sets this cookie to remember users' consent preferences so that their preferences are respected on subsequent visits to this site. It does not collect or store any personal information about the site visitors.\n\n\nFunctional\nFunctional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features.\n * Cookie\nlidc\n * Duration\n1 day\n * Description\nLinkedIn sets the lidc cookie to facilitate data center selection.\n\n\n * Cookie\nbrw\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nbrwConsent\n * Duration\n5 minutes\n * Description\nDescription is currently not available.\n\n\n * Cookie\nWMF-Uniq\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nloom_anon_comment\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nloom_referral_video\n * Duration\nsession\n * Description\nDescription is currently not available.\n\n\n * Cookie\nVISITOR_INFO1_LIVE\n * Duration\n6 months\n * Description\nA cookie set by YouTube to measure bandwidth that determines whether the user gets the new or old player interface.\n\n\n * Cookie\nyt-remote-connected-devices\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nytidb::LAST_RESULT_ENTRY_KEY\n * Duration\nNever Expires\n * Description\nThe cookie ytidb::LAST_RESULT_ENTRY_KEY is used by YouTube to store the last search result entry that was clicked by the user. This information is used to improve the user experience by providing more relevant search results in the future.\n\n\n * Cookie\nyt-remote-device-id\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-name\n * Duration\nsession\n * Description\nThe yt-remote-session-name cookie is used by YouTube to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\nyt-remote-fast-check-period\n * Duration\nsession\n * Description\nThe yt-remote-fast-check-period cookie is used by YouTube to store the user's video player preferences for embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-app\n * Duration\nsession\n * Description\nThe yt-remote-session-app cookie is used by YouTube to store user preferences and information about the interface of the embedded YouTube video player.\n\n\n * Cookie\nyt-remote-cast-available\n * Duration\nsession\n * Description\nThe yt-remote-cast-available cookie is used to store the user's preferences regarding whether casting is available on their YouTube video player.\n\n\n * Cookie\nyt-remote-cast-installed\n * Duration\nsession\n * Description\nThe yt-remote-cast-installed cookie is used to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\ncp_session\n * Duration\n3 months\n * Description\nCodepen sets this cookie for Help systems found in the website.\n\n\n * Cookie\nloid\n * Duration\n1 year 1 month 4 days\n * Description\nThis cookie is set by the Reddit. The cookie enables the sharing of content from the website onto the social media platform.\n\n\nAnalytics\nAnalytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc.\n * Cookie\n__hstc\n * Duration\n6 months\n * Description\nHubspot set this main cookie for tracking visitors. It contains the domain, initial timestamp (first visit), last timestamp (last visit), current timestamp (this visit), and session number (increments for each subsequent session).\n\n\n * Cookie\nhubspotutk\n * Duration\n6 months\n * Description\nHubSpot sets this cookie to keep track of the visitors to the website. This cookie is passed to HubSpot on form submission and used when deduplicating contacts.\n\n\n * Cookie\n_ga\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to calculate visitor, session and campaign data and track site usage for the site's analytics report. The cookie stores information anonymously and assigns a randomly generated number to recognise unique visitors.\n\n\n * Cookie\n_ga_*\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to store and count page views.\n\n\n * Cookie\n__Host-psifi.analyticsTrace\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n__Host-psifi.analyticsTraceV2\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n_gh_sess\n * Duration\nsession\n * Description\nGitHub sets this cookie for temporary application and framework state between pages like what step the user is on in a multiple step form.\n\n\n * Cookie\nYSC\n * Duration\nsession\n * Description\nYSC cookie is set by Youtube and is used to track the views of embedded videos on Youtube pages.\n\n\n * Cookie\najs_anonymous_id\n * Duration\n1 year\n * Description\nThis cookie is set by Segment to count the number of people who visit a certain site by tracking if they have visited before.\n\n\n * Cookie\nvuid\n * Duration\n1 year 1 month 4 days\n * Description\nVimeo installs this cookie to collect tracking information by setting a unique ID to embed videos on the website. \n\n\nPerformance\nPerformance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors.\n * Cookie\nAWSALB\n * Duration\n7 days\n * Description\nAWSALB is an application load balancer cookie set by Amazon Web Services to map the session to the target.\n\n\n * Cookie\nacq\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nacq.sig\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nptc\n * Duration\n2 years\n * Description\nNo description available.\n\n\nAdvertisement\nAdvertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns.\n * Cookie\nmuc_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to collect user behaviour and interaction data to optimize the website.\n\n\n * Cookie\nguest_id_marketing\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\nguest_id_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\npersonalization_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to integrate and share features for social media and also store information about how the user uses the website, for tracking and targeting.\n\n\n * Cookie\nguest_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor. It registers if a user is signed in to the Twitter platform and collects information about ad preferences.\n\n\n * Cookie\nbcookie\n * Duration\n1 year\n * Description\nLinkedIn sets this cookie from LinkedIn share buttons and ad tags to recognize browser IDs.\n\n\n * Cookie\n__Secure-ROLLOUT_TOKEN\n * Duration\n6 months\n * Description\nYouTube sets this cookie to manage feature rollout and experimentation. It helps Google control which new features or interface changes are shown to users as part of testing and staged rollouts, ensuring consistent experience for a given user during an experiment.\n\n\n * Cookie\nyt.innertube::nextId\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nyt.innertube::requests\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nsession_tracker\n * Duration\nsession\n * Description\nThis cookie is set by the Reddit. This cookie is used to identify trusted web traffic. It also helps in adverstising on the website.\n\n\n * Cookie\nedgebucket\n * Duration\nsession\n * Description\nReddit sets this cookie to save the information about a log-on Reddit user, for the purpose of advertisement recommendations and updating the content.\n\n\n * Cookie\ndid\n * Duration\n1 year\n * Description\nArbor sets this cookie to show targeted ads to site visitors.This cookie expires after 2 months or 1 year.\n\n\nUncategorised\nOther uncategorised cookies are those that are being analysed and have not been classified into a category as yet.\nNo cookies to display.\nReject All Save My Preferences Accept All\nToggle Search[](https://towardsdatascience.com/)Toggle Mobile Navigation\n[](https://towardsdatascience.com/)\nPublish AI, ML & data-science insights to a global community of data professionals.\nSign in[Submit an Article](https://write.insightmediagroup.io)\n[Latest](https://towardsdatascience.com/latest/)[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Newsletter](https://towardsdatascience.com/tag/the-variable/)[ShipAI](https://towardsdatascience.com/shipai/)\n* * *\n[Write For TDS](https://towardsdatascience.com/submissions/)\n[LinkedIn](https://www.linkedin.com/company/towards-data-science/)[X](https://x.com/TDataScience)Toggle Search\n[Data Engineering](https://towardsdatascience.com/category/data-science/data-engineering/)\n# Optimizing Data Transfer in Distributed AI/ML Training Workloads\nA deep dive on data transfer bottlenecks, their identification, and their resolution with the help of NVIDIA Nsight™ Systems – part 3\n[Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\nJanuary 23, 202613 min read\n![](../images/000.jpeg)\n[Photo](https://unsplash.com/photos/landscape-photography-of-river-with-trees-ssEQdOiKd8U) by Jachan DeVol on Unsplash\nThis is the third part of a series of posts on optimizing data transfer using NVIDIA Nsight™ Systems (nsys) profiler. [Part one](https://towardsdatascience.com/optimizing-data-transfer-in-ai-ml-workloads/) focused on CPU-to-GPU data copies, and [part two](https://towardsdatascience.com/optimizing-data-transfer-in-batched-ai-ml-inference-workloads/) on GPU-to-CPU copies. In this post, we turn our attention to data transfer between GPUs.\nNowadays, it is quite common for AI/ML training — particularly of large models — to be distributed across multiple GPUs. While there are many different schemes for performing such distribution, what they all have in common is their reliance on the constant transfer of data — such as gradients, weights, statistics, and/or metrics — between the GPUs, throughout training. As with the other types of data transfer we analyzed in our previous posts, here too, a poor implementation could easily lead to under-utilization of compute resources and the unjustified inflation of training costs. Optimizing GPU-to-GPU communication is an active area of research and innovation involving both hardware and software development.\nIn this post, we will focus on the most common form of distributed training — data-distributed training. In data-distributed training, identical copies of the ML model are maintained on each GPU. Each input batch is evenly distributed among the GPUs, each of which executes a training step to calculate the local gradients. The local gradients are then shared and averaged across the GPUs, resulting in an identical gradient update to each of the model copies. Using [NVIDIA Nsight™ Systems (nsys)](https://docs.nvidia.com/nsight-systems/index.html) profiler, we will analyze the effect of the GPU-to-GPU transfer of the gradients on the runtime performance of training a toy model and assess a few techniques for reducing its overhead.\nLearn this step by step with the interactive [Data Engineer](https://roadmap.sh/data-engineer) roadmap.\n#### Disclaimers\nThe code we will share in this post is intended for demonstrative purposes; please do not rely on its accuracy or optimality. Please do not interpret our mention of any tool, framework, library, service, or platform as an endorsement of its use.\nThanks to [Yitzhak Levi](https://www.linkedin.com/in/yitzhak-levi-49a217201/) for his contributions to this post.\n## Instance Selection for Distributed Training\nIn a previous post, [Instance Selection for Deep Learning](https://towardsdatascience.com/instance-selection-for-deep-learning-7463d774cff0/), we discussed the importance of choosing an instance type that is most suitable for your AI/ML workload and the potential impact of that choice on the success of your project. When choosing an instance type for a workload that includes a lot of GPU-to-GPU traffic, you will want to pay attention to how such communication is carried out, including: the instance topology, GPU interconnects, maximal throughput, and latency.\nIn this post, we will limit our discussion to distribution between GPUs on a single instance. We will experiment with two instance types: an [Amazon EC2 g6e.48xlarge](https://aws.amazon.com/ec2/instance-types/g6e/) with 8 [NVIDIA L40S](https://www.nvidia.com/en-eu/data-center/l40s/) GPUs and an [Amazon EC2 p4d.24xlarge](https://aws.amazon.com/ec2/instance-types/p4/) with 8 [NVIDIA A100](https://www.nvidia.com/en-eu/data-center/a100/) GPUs. Each will run an [AWS Deep Learning (Ubuntu 24.04) AMI](https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html) with [PyTorch (2.8)](https://pytorch.org/get-started/locally/), [nsys-cli profiler](https://developer.nvidia.com/nsight-systems/get-started) (version 2025.6.1), and the [NVIDIA Tools Extension](https://docs.nvidia.com/tools-extension/index.html) (NVTX) library.\nOne of the primary differences between the two instance types is how the GPUs are connected: On the [g6e.48xlarge](https://aws.amazon.com/ec2/instance-types/g6e/) communication between the GPUs is over [PCI Express (PCIe)](https://en.wikipedia.org/wiki/PCI_Express), whereas [p4d.24xlarge](https://aws.amazon.com/ec2/instance-types/p4/) includes [NVIDIA NVLink™](https://www.nvidia.com/en-eu/data-center/nvlink/) — dedicated hardware for enabling high-throughput GPU-to-GPU communication. Communication over the PCIe bus is significantly slower than NVLink. While this may be sufficient for workloads with a low communication-to-compute ratio, it could be a performance killer for workloads with high communication rates.\nTo discover the topology of the instance types, we run the following command:\nbash\n\n```\nnvidia-smi topo -m\n```\n\nOn the [g6e.48xlarge](https://aws.amazon.com/ec2/instance-types/g6e/), we get the following results:\n![](../images/001.webp)\nGPU Topology of [g6e.48xlarge](https://aws.amazon.com/ec2/instance-types/g6e/) (by Author)\nGPUs on the same NUMA node are connected by a “NODE” link and GPUs on different NUMA nodes by a “SYS” link. Both links traverse the PCIe as well as another interconnect; neither is a direct connection (a.k.a., a “PIX” link). We will see later on how this can impact throughput performance.\nOn the [p4d.24xlarge](https://aws.amazon.com/ec2/instance-types/p4/), every pair of GPUs is linked by a dedicated NVLink connection:\n![](../images/body-29c24f88d241224f788e.webp)\nGPU Topology of [p4d.24xlarge](https://aws.amazon.com/ec2/instance-types/p4/) (by Author)\n## A Toy Model\nTo facilitate our discussion, we build a toy data-distributed training experiment.\nWe choose a model with a relatively high communication-to-compute ratio — a [Vision Transformer (ViT)-L/32](https://docs.pytorch.org/vision/main/models/generated/torchvision.models.vit_l_32.html#torchvision.models.vit_l_32) image-classification model with roughly 306 million parameters — and a synthetic dataset that we will use to train it:\npython\n\n```\nimport os, time, torch, nvtx\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.distributed as dist\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision.models import vit_l_32\n\nWORLD_SIZE = int(os.environ.get(\"WORLD_SIZE\", 1))\nBATCH_SIZE = 32\nIMG_SIZE = 224\nWARMUP_STEPS = 10\nPROFILE_STEPS = 3\nCOOLDOWN_STEPS = 1\nTOTAL_STEPS = WARMUP_STEPS + PROFILE_STEPS + COOLDOWN_STEPS\nN_WORKERS = 8\n\ndef get_model():\n return vit_l_32(weights=None)\n\n# A synthetic dataset with random images and labels\nclass FakeDataset(Dataset):\n def __len__(self):\n return TOTAL_STEPS * BATCH_SIZE * WORLD_SIZE\n\n def __getitem__(self, index):\n img = torch.randn((3, IMG_SIZE, IMG_SIZE))\n label = torch.randint(0, 1000, (1,)).item()\n return img, label\n\ndef get_data_iter(rank, world_size):\n dataset = FakeDataset()\n\n sampler = DistributedSampler(dataset, num_replicas=world_size,\n rank=rank, shuffle=True)\n\n train_loader = DataLoader(dataset, batch_size=BATCH_SIZE,\n sampler=sampler, num_workers=N_WORKERS, \n pin_memory=True)\n\n return iter(train_loader)\n```\n\nWe define a utility function that we will use to set up PyTorch [DistributedDataParallel](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) (DDP) training. We configure the PyTorch [distributed communication package](https://docs.pytorch.org/docs/stable/distributed.html) to use the [NVIDIA Collective Communications Library](https://developer.nvidia.com/nccl) (NCCL) as its [communication backend](https://docs.pytorch.org/tutorials/intermediate/dist_tuto.html#communication-backends) and wrap the model in a [DistributedDataParallel](https://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) container.\npython\n\n```\ndef configure_ddp(model, rank):\n dist.init_process_group(\"nccl\")\n model = DDP(model,\n device_ids=[rank],\n bucket_cap_mb=2000)\n return model\n```\n\nNote that we naively configure the DDP bucket capacity to 2 GB. DDP groups the model’s gradients into buckets and performs gradient reduction of each bucket in a separate command. The implication of setting the bucket capacity to 2 GB is that all of the ~306 million (FP32) gradients will fit in a single bucket (306 million x 4 bytes per gradient = ~1.22 GB) and reduction will only occur once all gradients have been calculated.\nAs in our previous posts, we schedule nsys profiler [programmatically](https://github.com/pytorch/pytorch/blob/v2.9.0/torch/cuda/profiler.py#L50) and wrap the different portions of the training step with color-coded [NVTX](https://nvtx.readthedocs.io/en/latest/index.html) annotations:\npython\n\n```\ndef train(use_ddp=False):\n # detect the env vars set by torchrun\n rank = int(os.environ.get(\"RANK\", 0))\n local_rank = int(os.environ.get(\"LOCAL_RANK\", 0))\n torch.cuda.set_device(local_rank)\n\n model = get_model().to(local_rank)\n criterion = nn.CrossEntropyLoss().to(local_rank)\n\n if use_ddp:\n model = configure_ddp(model, rank)\n\n optimizer = optim.SGD(model.parameters())\n\n data_iter = get_data_iter(rank, WORLD_SIZE)\n\n model.train()\n\n for i in range(TOTAL_STEPS):\n\n # Schedule Profiling\n if i == WARMUP_STEPS:\n torch.cuda.synchronize()\n start_time = time.perf_counter()\n torch.cuda.profiler.start()\n elif i == WARMUP_STEPS + PROFILE_STEPS:\n torch.cuda.synchronize()\n torch.cuda.profiler.stop()\n end_time = time.perf_counter()\n\n with nvtx.annotate(f\"Batch {i}\", color=\"blue\"):\n\n with nvtx.annotate(\"get batch\", color=\"red\"):\n data, target = next(data_iter)\n data = data.to(local_rank, non_blocking=True)\n target = target.to(local_rank, non_blocking=True)\n with nvtx.annotate(\"forward\", color=\"green\"):\n output = model(data)\n loss = criterion(output, target)\n\n with nvtx.annotate(\"backward\", color=\"purple\"):\n optimizer.zero_grad()\n loss.backward()\n\n with nvtx.annotate(\"optimizer step\", color=\"yellow\"):\n optimizer.step()\n\n if use_ddp:\n dist.destroy_process_group()\n\n if rank == 0:\n total_time = end_time - start_time\n print(f\"Throughput: {PROFILE_STEPS/total_time:.2f} steps/sec\")\n\n\nif __name__ == \"__main__\":\n # enable ddp if run with torchrun\n train(use_ddp=\"RANK\" in os.environ)\n```\n\nWe set the `NCCL_DEBUG` environment variable for visibility into how NCCL sets up the transport links between the GPUs.\nbash\n\n```\nexport NCCL_DEBUG=INFO\n```\n\n### Single-GPU Performance\nWe begin our experimentation by running our script on a single GPU without DDP:\nbash\n\n```\nnsys profile \\\n --capture-range=cudaProfilerApi \\\n --capture-range-end=stop \\\n --trace=cuda,nvtx,osrt,nccl \\\n --output=baseline \\\n python train.py\n```\n\nNote the inclusion of `nccl`in the _trace_ section; this will be critical for analyzing the GPU-to-GPU communication in the multi-GPU experiments.\nThe throughput of the baseline experiment is 8.91 steps per second on the L40S GPU and 5.17 steps per second on the A100: The newer NVIDIA Ada Lovelace Architecture performs better than the NVIDIA Ampere Architecture on our toy model. In the image below we show the timeline view for the L40S experiment.\n![](../images/body-2641afc81297bc88560d.webp)\nSingle L40S GPU Nsight Systems Trace (by Author)\nIn this post we will focus on the CUDA portion of the timeline. Looking at the NVTX row, we see the recurring red (CPU-to-GPU copy), green (forward), purple (backward), and yellow (optimizer update) pattern that makes up our train step. Note, however, that the purple appears as just a tiny blip while, in practice, the backward pass fills the entire gap between the forward and optimizer blocks: The NVTX library appears to have captured only the initial launch of the autograd graph.\nWe will use this trace as a comparative baseline for our next experiments.\n### DDP with 1 GPU\nWe assess the impact of the DDP wrapper by running [torchrun](https://docs.pytorch.org/docs/stable/elastic/run.html) on a single GPU:\nbash\n\n```\ntorchrun --nproc_per_node=1 \\\n --no-python \\\n nsys profile \\\n --capture-range=cudaProfilerApi \\\n --capture-range-end=stop \\\n --trace=cuda,nvtx,nccl,osrt \\\n --output=ddp-1gpu \\\n python train.py\n```\n\nThe resulting throughput drops to 8.40 steps per second on the L40S GPU and 5.04 steps per second on the A100. Even in the absence of any cross-GPU communication, DDP introduces overhead that can decrease throughput by ~3–7%. The important lesson from this is to always ensure that single GPU training is run without DDP.\nThe nsys trace of the L40S experiment confirms the presence of the DDP overhead:\n![](../images/body-8a0671e4ae02e65e70b7.webp)\nSingle L40S GPU With DDP - Profiler Trace (by Author)\nThe main change from the training step in the previous trace is a large chunk of device-to-device (_DtoD_) memory copies at the end of the backward block and just before the optimizer block (highlighted in the trace above). This is the DDP in action: Even in the absence of a cross GPU gradient reduction, DDP prepares the local gradients in a dedicated memory block for reduction and then copies the results back into the grad property of each parameter in preparation for the gradient update. (Note that the _DtoD_ copies are between two memory locations on the same GPU — not between two different GPUs).\n### DDP With Multiple GPUs\nNext, we assess the impact of gradient sharing between 2, 4, and 8 GPUs:\nbash\n\n```\ntorchrun --nproc_per_node=8 \\\n --no-python \\\n nsys profile \\\n --capture-range=cudaProfilerApi \\\n --capture-range-end=stop \\\n --trace=cuda,nvtx,nccl,osrt \\\n --output=ddp-8gpu_%q{RANK} \\\n python train.py\n```\n\nThe table below captures impact on the training throughput on the [g6e.48xlarge](https://aws.amazon.com/ec2/instance-types/g6e/) and the [p4d.24xlarge](https://aws.amazon.com/ec2/instance-types/p4/):\n![](../images/body-b684d1a624dc89d32e13.webp)\nThroughput Performance of DDP Training (by Author)\nDDP training performance plummets on the g6e.48xlarge instance: The 8-GPU throughput is more than 6 times slower than the single-GPU result. The NCCL logs include multiple lines describing the communication paths between the GPUs, e.g.:\nbash\n\n```\nNCCL INFO Channel 00 : 0[0] -> 1[1] via SHM/direct/direct\n```\n\nThis indicates that the data transfer between each two GPUs passes through CPU shared memory which greatly limits the communication bandwidth.\nOn the p4d.24xlarge, in contrast, where the NVLink connections allow for direct peer-to-peer (P2P) communication, the 8-GPU throughput is just 8% lower than the baseline result:\nbash\n\n```\nNCCL INFO Channel 00/0 : 0[0] -> 1[1] via P2P/CUMEM/read\n```\n\nEven though each individual L40S outperforms the A100 on our toy model by 72%, the inclusion of NVLink makes the p4d.24xlarge more optimal than the g6e.48xlarge for running DDP over 8 GPUs.\nThe data transfer bottleneck on the g6e.48xlarge can be easily seen in the nsys trace. Here we use the “multiple view” option to display the activity of all the DDP processes in a single timeline:\n![](../images/body-7295100c521095e6f8eb.webp)\n8-GPU DDP on g6e.48xlarge - Profiler Trace (by Author)\nThe gradient reduction occurs in the _ncclAllReduce_ call on the NCCL row. This comes just after the local gradient calculation has completed and just before the _DtoD_ memory operation discussed above that copies the reduced gradients back to the _grad_ field of each parameter. On the g6e.48xlarge the NCCL operation dominates the backward pass, accounting for approximately 84% of the overall step time (587 out of 701 milliseconds).\nOn the p4d.24xlarge, the NCCL call takes up a much smaller portion of the training step:\n![](../images/body-e34e43c7563d7c3cf0b9.webp)\n8-GPU DDP on p4d.24xlarge - Profiler Trace (by Author)\nIn the next sections we will discuss a few DDP optimization techniques and assess their impact on runtime performance. We will limit our scope to PyTorch-level optimizations, leaving other methods (e.g., NCCL/OS/network tuning) for another post.\nA general approach for dealing with communication bottlenecks is to reduce the frequency of communication. In DDP workloads this can be achieved by increasing the data batch size (i.e., increasing the compute per step) or, if the memory capacity forbids this, apply gradient accumulation — instead of applying the gradient reduction and update every step, accumulate the local gradients for _N_ steps and apply the reduction every _N_ th step. Both techniques increase the _effective_ batch size — the number of overall samples between gradient updates which can impact model convergence and may require tuning of the optimization parameters. In our discussion we assume the _effective_ batch size is fixed.\nWe will cover four techniques. The first two focus on reducing DDP overhead. The final two directly address the data transfer bottleneck.\n## Optimization 1: Static Graph Declaration\nOur first change is to pass `static_graph=True` to the [DDP](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) container. Generally speaking, models might include parameters for which gradients are not calculated at every step — referred to as “unused parameters” in the DDP documentation. This is common in models with conditional logic. DDP includes dedicated mechanisms for identifying and handling unused parameters. In the case of our toy model, all of the gradients are calculated in each step — our graph is “static”. Explicitly declaring our graph as static reduces the overhead associated with handling a dynamic gradient set.\npython\n\n```\ndef configure_ddp(model, rank):\n dist.init_process_group(\"nccl\")\n model = DDP(model,\n device_ids=[rank],\n static_graph=True,\n bucket_cap_mb=2000)\n return model\n```\n\nIn the case of our toy model, the impact of this change is negligible on both the g6e.48xlarge and the p4d.24xlarge. Without any further delay, we proceed to the next optimization.\n## Optimization 2: Increase Memory Efficiency\nOur second technique addresses the large chunk of _DtoD_ memory copying we identified above. Instead of copying the gradients to and from the NCCL communication memory blocks, we can explicitly set the parameter gradients to point directly to the NCCL communication buffers. Consequently, the same memory is used to store the local gradients, to perform the gradient reduction, and to apply the gradient updates. This is configured by passing `gradient_as_bucket_view=True` to the [DDP](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) container:\npython\n\n```\ndef configure_ddp(model, rank):\n dist.init_process_group(\"nccl\")\n model = DDP(model,\n device_ids=[rank],\n static_graph=True,\n gradient_as_bucket_view=True,\n bucket_cap_mb=2000)\n return model\n```\n\nIn the trace below, captured on the p4d.24xlarge, we no longer see the block of _DtoD_ memory copy between the all-reduce and (yellow) optimizer steps:\n![](../images/body-e2b4f30866a41d538570.webp)\nMemory-Optimized DDP on p4d.24xlarge - Profiler Trace (by Author)\nIn the case of our toy example, this optimization boosts performance by a modest 1%.\n## Optimization 3: Gradient Compression\nA common approach for addressing communication bottlenecks is to apply compression algorithms to reduce the size of the payload. PyTorch DDP provides a number of [dedicated communication hooks](https://docs.pytorch.org/docs/stable/ddp_comm_hooks.html) that automate compression of gradients before NCCL reduction and decompression afterward. In the code block below, we apply [bfloat16 gradient compression](https://docs.pytorch.org/docs/stable/ddp_comm_hooks.html#torch.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_hook):\npython\n\n```\nimport torch.distributed.algorithms.ddp_comm_hooks.default_hooks as ddp_hks\n\ndef configure_ddp(model, rank):\n dist.init_process_group(\"nccl\")\n model = DDP(model,\n device_ids=[rank],\n static_graph=True,\n gradient_as_bucket_view=True,\n bucket_cap_mb=2000)\n\n model.register_comm_hook(state=None, hook=ddp_hks.bf16_compress_hook)\n return model\n```\n\nOn the heavily bottlenecked g6e.48xlarge instance bfloat16 compression results in a considerable 65% speed-up! The nsys trace shows the reduced-sized NCCL call as well as the newly introduced compression operations:\n![](../images/body-31ed178066044ddb30bf.webp)\nBF16-Compressed DDP on g6e.48xlarge - Profiler Trace (by Author)\nOn the p4d.24xlarge the overhead of the compression operations outweigh the gains in the communication speed, leading to an overall reduction in throughput:\n![](../images/body-3cdbf39e9f0cc5feb383.webp)\nBF16-Compressed DDP on p4d.24xlarge - Profiler Trace (by Author)\nPyTorch offers a more aggressive compression algorithm than bfloat16, called [PowerSGD](https://docs.pytorch.org/docs/stable/ddp_comm_hooks.html#powersgd-communication-hook). Below we present a naive usage — in practice this can require a lot of tuning. Please see the documentation for details:\npython\n\n```\nfrom torch.distributed.algorithms.ddp_comm_hooks import powerSGD_hook \n\ndef configure_ddp(model, rank):\n dist.init_process_group(\"nccl\")\n model = DDP(model,\n device_ids=[rank],\n static_graph=True,\n gradient_as_bucket_view=True,\n bucket_cap_mb=2000)\n\n state = powerSGD_hook.PowerSGDState(\n process_group=None\n )\n\n model.register_comm_hook(\n state,\n hook=powerSGD_hook.powerSGD_hook\n )\n return model\n```\n\n[PowerSGD](https://docs.pytorch.org/docs/stable/ddp_comm_hooks.html#powersgd-communication-hook) has a dramatic impact on the g6e.48xlarge instance, increasing throughput all the way back up to 7.5 steps per second — more than five times faster than the baseline result! Note the reduction in NCCL kernel size in the resultant trace:\n![](../images/body-0154b97c75d547462e1a.webp)\nPowerSGD-Compressed DDP on g6e.48xlarge - Profiler Trace (by Author)\nIt is important to note that these compression algorithms are precision-lossy and should be used with caution, as they may impact your model convergence.\n## Optimization 4: Parallelize Gradient Reduction\nDDP groups parameters into multiple buckets and triggers gradient reduction of each bucket independently — as soon as the bucket is filled. This allows for running gradient reduction of filled buckets while gradient calculation (of other buckets) is still ongoing. The degree of overlap depends on the number of DDP buckets which we control via the _bucket_cap_mb_ setting of the DDP container. Recall that in our initial implementation, we explicitly set this to _2000_(2 GB) which (given the size of our model) translated to a single DDP bucket. The table below shows the throughput for different values of _bucket_cap_mb_. The optimal setting will vary based on the details of the model and runtime environment.\n![](../images/body-36d6588037899af165d5.webp)\nDDP Throughput vs. Bucket Capacity (by Author)\nNote the significant 12% improvement when using multiple buckets on the g6e.48xlarge with BF16 compression. PowerSGD, on the other hand, works best when applied to a single bucket.\nIn the image below, captured on a p4d.24xlarge with _bucket_cap_mb_ set to _100_ , we can see the impact of this optimization on the profiler trace:\n![](../images/body-da09f046e26c5cee836a.webp)\nDDP With Multiple Buckets on p4d.24xlarge - Profiler Trace (by Author)\nIn place of the single NCCL all-reduce operation, we now have 11 smaller blocks running in parallel with the local gradient computation.\n## Results\nWe summarize the results of our experiments in the following table:\n![](../images/body-7665aa46ed637eed725d.webp)\nExperiment Results (by Author)\nOn the p4d.24xlarge, where the data is transferred over NVLink, the overall impact was a modest 4% speed-up. But on the g6e.48xlarge the gains were significant and mostly due to gradient compression — and 86% boost for the BF16 scheme and an over 5X improvement for (the naive implementation of) PowerSGD.\nImportantly, these results can vary considerably based on the model architecture and runtime environment.\n## Summary\nThis concludes our three-part series on identifying and solving data transfer bottlenecks with NVIDIA Nsight™ Systems (nsys) profiler. In each of the posts we demonstrated the tendency of data transfers to introduce performance bottlenecks and resource under-utilization. In each case, we used nsys profiler to identify the bottlenecks and measure the impact of different optimization techniques.\nThe accelerations we achieved in each of the use-cases that we studied strengthen the importance of integrating the regular use of tools such as nsys profiler into the AI/ML development workflow and highlight the opportunity — even for non-CUDA specialists — to achieve meaningful performance gains and AI/ML cost reductions.\n## Register Now to Continue Reading\n### Unlock Towards Data Science with a Free Account\nAll articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more.\nRegister for freeSign in to your account\n* * *\n![](../images/body-97f5483d57a9b8741146.jpeg)\nWritten By\n## [Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\n[See all from Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\n[Deep Dives](https://towardsdatascience.com/tag/deep-dives/)[Artificial Intelligence](https://towardsdatascience.com/tag/artificial-intelligence/)[Deep Learning](https://towardsdatascience.com/tag/deep-learning/)[Nvidia Gpu](https://towardsdatascience.com/tag/nvidia-gpu/)[Pytorch](https://towardsdatascience.com/tag/pytorch/)\nShare This Article\n * [](https://www.facebook.com/sharer/sharer.php?u=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-data-transfer-in-distributed-ai-ml-training-workloads%2F&title=Optimizing%20Data%20Transfer%20in%20Distributed%20AI%2FML%20Training%20Workloads)\n * [](https://www.linkedin.com/shareArticle?mini=true&url=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-data-transfer-in-distributed-ai-ml-training-workloads%2F&title=Optimizing%20Data%20Transfer%20in%20Distributed%20AI%2FML%20Training%20Workloads)\n * [](https://x.com/share?url=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-data-transfer-in-distributed-ai-ml-training-workloads%2F&text=Optimizing%20Data%20Transfer%20in%20Distributed%20AI%2FML%20Training%20Workloads)\n\n\nTowards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program.\n[Write for TDS](https://towardsdatascience.com/submissions/)\n## Related Articles\n * [![](../images/002.jpeg)](https://towardsdatascience.com/all-data-integrations-should-use-change-data-capture-a1d207091773/)\n## [All Data Integrations Should Use Change Data Capture](https://towardsdatascience.com/all-data-integrations-should-use-change-data-capture-a1d207091773/)\n[Data Engineering](https://towardsdatascience.com/category/data-science/data-engineering/)\nCDC elevates the game for data integrations. But it is hard to to get right - learn what it takes to build data integrations with CDC.\n[Raghotham Murthy](https://towardsdatascience.com/author/raghotham/)\nApril 21, 202111 min read\n * [![](../images/003.jpeg)](https://towardsdatascience.com/dbt-models-structure-c31c8977b5fc/)\n## [How to Structure Your dbt Projects and Data Models](https://towardsdatascience.com/dbt-models-structure-c31c8977b5fc/)\n[Programming](https://towardsdatascience.com/category/programming/)\nEnforcing meaningful structure to dbt data models\n[Giorgos Myrianthous](https://towardsdatascience.com/author/gmyrianthous/)\nNovember 9, 20227 min read\n * [![](../images/004.jpeg)](https://towardsdatascience.com/10-most-practical-data-science-skills-you-should-know-in-2022-9487d7750e8a/)\n## [10 Most Practical Data Science Skills You Should Know in 2022](https://towardsdatascience.com/10-most-practical-data-science-skills-you-should-know-in-2022-9487d7750e8a/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nSkills that will actually make you employable\n[Terence Shin, MSc, MBA](https://towardsdatascience.com/author/terenceshin/)\nSeptember 5, 20217 min read\n * [![](../images/005.png)](https://towardsdatascience.com/how-to-identify-your-business-critical-data-c58506116a41/)\n## [How to identify your business-critical data](https://towardsdatascience.com/how-to-identify-your-business-critical-data-c58506116a41/)\n[Data Science](https://towardsdatascience.com/category/data-science/)\nPractical steps to identifying business-critical data models and dashboards and drive confidence in your data\n[Mikkel Dengsøe](https://towardsdatascience.com/author/mikldd/)\nJune 14, 20239 min read\n * [![](../images/006.jpeg)](https://towardsdatascience.com/to-mask-or-not-to-mask-the-effect-of-prompt-tokens-on-instruction-tuning-016f85fd67f4/)\n## [To Mask or Not to Mask: The Effect of Prompt Tokens on Instruction Tuning](https://towardsdatascience.com/to-mask-or-not-to-mask-the-effect-of-prompt-tokens-on-instruction-tuning-016f85fd67f4/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nImplementing prompt-loss-weight, and why we should replace prompt-masking with prompt-weighting\n[David Vaughn](https://towardsdatascience.com/author/davidsvaughn/)\nSeptember 30, 202431 min read\n * [![](../images/007.jpg)](https://towardsdatascience.com/7-ways-to-improve-performance-of-your-data-warehouse-4544456bcc3/)\n## [7 Ways to Improve Performance of Your Data Warehouse](https://towardsdatascience.com/7-ways-to-improve-performance-of-your-data-warehouse-4544456bcc3/)\n[Data Science](https://towardsdatascience.com/category/data-science/)\nImprove Poor Performance of Your Data Warehouse by Applying These Steps\n[Hanzala Qureshi](https://towardsdatascience.com/author/hanzalaqureshi/)\nApril 26, 20226 min read\n * [![](../images/008.jpg)](https://towardsdatascience.com/easy-parquet-tutorial-best-practices-237955e46cb7/)\n## [Simple Parquet Tutorial and Best Practices](https://towardsdatascience.com/easy-parquet-tutorial-best-practices-237955e46cb7/)\n[Data Science](https://towardsdatascience.com/category/data-science/)\nHands-on tutorial for starting your Parquet learning\n[Arli](https://towardsdatascience.com/author/arli94/)\nDecember 16, 20228 min read\n * [![](../images/009.jpeg)](https://towardsdatascience.com/unit-testing-pyspark-code-using-pytest-b5ab2fd54415/)\n## [Unit testing PySpark code using Pytest](https://towardsdatascience.com/unit-testing-pyspark-code-using-pytest-b5ab2fd54415/)\n[Data Science](https://towardsdatascience.com/category/data-science/)\nWhen it comes to writing unit-tests for PySpark pipelines, writing focussed, fast, isolated and concise tests can be a challenge.\n[Julian West](https://towardsdatascience.com/author/julianwest155/)\nJanuary 16, 202311 min read\n * [![](../images/010.png)](https://towardsdatascience.com/structured-state-space-models-visually-explained-86cfe2757386/)\n## [Structured State Space Models Visually Explained](https://towardsdatascience.com/structured-state-space-models-visually-explained-86cfe2757386/)\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\nPart 2 - Towards Mamba State Space Models for Images, Videos and Time Series\n[Sascha Kirch](https://towardsdatascience.com/author/saschakirch/)\nAugust 22, 202418 min read\n\n\n * [YouTube](https://www.youtube.com/c/TowardsDataScience)\n * [X](https://x.com/TDataScience)\n * [LinkedIn](https://www.linkedin.com/company/towards-data-science/)\n * [Threads](https://www.threads.net/@towardsdatascience)\n * [Bluesky](https://bsky.app/profile/towardsdatascience.com)\n\n\n[](https://towardsdatascience.com/)\nYour home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals.\nSubscribe to Our Newsletter[Write For TDS](https://towardsdatascience.com/submissions/)•[About](https://towardsdatascience.com/about-towards-data-science-d691af11cc2f/)•[Advertise](https://contact.towardsdatascience.com/advertise-with-towards-data-science)•[Privacy Policy](https://towardsdatascience.com/privacy-policy/)•[Terms of Use](https://towardsdatascience.com/website-terms-of-use/)\n© Insight Media Group, LLC 2026•[The New Stack](https://thenewstack.io)•[roadmap.sh](https://roadmap.sh)•[Creator Network](https://insightmediagroup.io/creator-network)\n#### Recommended Articles \nClose\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500) Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1 ](https://towardsdatascience.com/javascript-fatigue-you-dont-need-js-to-build-chatgpt/ \"Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500) The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs ](https://towardsdatascience.com/subset-sum-problem-solved-in-linear-time-for-dense-enough-inputs/ \"The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500) Does Your Company Have a Data Strategy? ](https://towardsdatascience.com/does-your-company-have-a-data-strategy-e8944dbec4ae/ \"Does Your Company Have a Data Strategy?\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500) Productivity Tips, Data Career Insights, and Other Recent Must-Reads ](https://towardsdatascience.com/productivity-tips-data-career-insights-and-other-recent-must-reads-5c47eb7a2fe0/ \"Productivity Tips, Data Career Insights, and Other Recent Must-Reads\")\n\n\n", "metadata": { "title": "Optimizing Data Transfer in Distributed AI/ML Training Workloads | Towards Data Science", "description": "A deep dive on data transfer bottlenecks, their identification, and their resolution with the help of NVIDIA Nsight™ Systems – part 3", "keywords": null, "author": "Chaim Rand", "og:locale": "en_US", "og:type": "article", "og:title": "Optimizing Data Transfer in Distributed AI/ML Training Workloads", "og:description": "A deep dive on data transfer bottlenecks, their identification, and their resolution with the help of NVIDIA Nsight™ Systems – part 3", "og:url": "https://towardsdatascience.com/optimizing-data-transfer-in-distributed-ai-ml-training-workloads/", "og:site_name": "Towards Data Science", "og:image": "https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/1-2-1.jpeg", "og:image:width": "1100", "og:image:height": "726", "og:image:type": "image/jpeg", "twitter:card": "summary_large_image", "twitter:site": "@TDataScience", "twitter:creator": "@TDataScience", "twitter:label1": "Written by", "twitter:data1": "Chaim Rand", "twitter:label2": "Est. reading time", "twitter:data2": "13 min read", "article:published_time": "2026-01-23T11:30:09.000Z", "article:modified_time": "2026-01-23T11:30:09.000Z", "article:author": "https://towardsdatascience.com/author/chaimrand/" }, "paywall": { "listed_paid": false, "detected": false, "marker": "", "public_only": false, "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙" }, "crawl": { "success": true, "status_code": 202, "error": "" }, "code_blocks": [ { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_000.txt", "chars": 19 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_001.txt", "chars": 1310 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_002.txt", "chars": 179 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_003.txt", "chars": 1813 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_004.txt", "chars": 23 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_005.txt", "chars": 152 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_006.txt", "chars": 214 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_007.txt", "chars": 223 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_008.txt", "chars": 58 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_009.txt", "chars": 57 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_010.txt", "chars": 214 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_011.txt", "chars": 260 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_012.txt", "chars": 412 }, { "language": "txt", "file": "attachments/f29ca4fd35ae/code/block_013.txt", "chars": 500 } ], "images": [ { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T18:06:18Z" }, { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T18:06:18Z" }, { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T18:06:18Z" }, { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T18:06:18Z" } ], "audio": [], "videos": [ { "url": "https://www.youtube.com/c/TowardsDataScience", "kind": "channel_reference", "status": "reference_only" } ], "links": [ { "type": "internal", "href": "https://towardsdatascience.com/" }, { "type": "internal", "href": "https://towardsdatascience.com/latest/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/editors-pick/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/the-variable/" }, { "type": "internal", "href": "https://towardsdatascience.com/shipai/" }, { "type": "internal", "href": "https://towardsdatascience.com/submissions/" }, { "type": "internal", "href": "https://towardsdatascience.com/category/data-science/data-engineering/" }, { "type": "internal", "href": "https://towardsdatascience.com/author/chaimrand/" }, { "type": "internal", "href": "https://towardsdatascience.com/optimizing-data-transfer-in-ai-ml-workloads/" }, { "type": "internal", "href": "https://towardsdatascience.com/optimizing-data-transfer-in-batched-ai-ml-inference-workloads/" }, { "type": "internal", "href": "https://towardsdatascience.com/instance-selection-for-deep-learning-7463d774cff0/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/deep-dives/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/artificial-intelligence/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/deep-learning/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/nvidia-gpu/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/pytorch/" }, { "type": "internal", "href": "https://towardsdatascience.com/all-data-integrations-should-use-change-data-capture-a1d207091773/" }, { "type": "internal", "href": "https://towardsdatascience.com/author/raghotham/" }, { "type": "internal", "href": "https://towardsdatascience.com/dbt-models-structure-c31c8977b5fc/" }, { "type": "internal", "href": "https://towardsdatascience.com/category/programming/" }, { "type": "external", "href": "https://write.insightmediagroup.io" }, { "type": "external", "href": "https://www.linkedin.com/company/towards-data-science/" }, { "type": "external", "href": "https://x.com/TDataScience" }, { "type": "external", "href": "https://unsplash.com/photos/landscape-photography-of-river-with-trees-ssEQdOiKd8U" }, { "type": "external", "href": "https://docs.nvidia.com/nsight-systems/index.html" }, { "type": "external", "href": "https://roadmap.sh/data-engineer" }, { "type": "external", "href": "https://www.linkedin.com/in/yitzhak-levi-49a217201/" }, { "type": "external", "href": "https://aws.amazon.com/ec2/instance-types/g6e/" }, { "type": "external", "href": "https://www.nvidia.com/en-eu/data-center/l40s/" }, { "type": "external", "href": "https://aws.amazon.com/ec2/instance-types/p4/" }, { "type": "external", "href": "https://www.nvidia.com/en-eu/data-center/a100/" }, { "type": "external", "href": "https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html" }, { "type": "external", "href": "https://pytorch.org/get-started/locally/" }, { "type": "external", "href": "https://developer.nvidia.com/nsight-systems/get-started" }, { "type": "external", "href": "https://docs.nvidia.com/tools-extension/index.html" }, { "type": "external", "href": "https://en.wikipedia.org/wiki/PCI_Express" }, { "type": "external", "href": "https://www.nvidia.com/en-eu/data-center/nvlink/" }, { "type": "external", "href": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.vit_l_32.html" }, { "type": "external", "href": "https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html" }, { "type": "external", "href": "https://docs.pytorch.org/docs/stable/distributed.html" } ], "repair_version": 2, "repaired_at": "2026-10-20T17:32:37Z", "source_url": "https://towardsdatascience.com/optimizing-data-transfer-in-distributed-ai-ml-training-workloads/", "delivery_status": "unverified", "discovery_status": "http_403", "missing_reason": "源页面未重新核验:http_403；已按保留正文补采图片", "repair_error": "", "actual_video_files": [], "expected_saved_body_images": 4, "image_discovery": "saved_markdown", "source_body_markdown": "![Revisit consent button](https://cdn-cookieyes.com/assets/images/revisit.svg)\nWe value your privacy\nWe use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking \"Accept All\", you consent to our use of cookies.\nCustomise Reject All Accept All\nCustomise Consent Preferences ![](https://cdn-cookieyes.com/assets/images/close.svg)\nWe use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below.\nThe cookies that are categorised as \"Necessary\" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more\nNecessaryAlways Active\nNecessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data.\n * Cookie\nBCTempID\n * Duration\n10 minutes\n * Description\nNo description available.\n\n\n * Cookie\n__cf_bm\n * Duration\n1 hour\n * Description\nThis cookie, set by Cloudflare, is used to support Cloudflare Bot Management. \n\n\n * Cookie\nAWSALBCORS\n * Duration\n7 days\n * Description\nAmazon Web Services set this cookie for load balancing.\n\n\n * Cookie\n_cfuvid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to track users across sessions to optimize user experience by maintaining session consistency and providing personalized services\n\n\n * Cookie\nli_gc\n * Duration\n6 months\n * Description\nLinkedin set this cookie for storing visitor's consent regarding using cookies for non-essential purposes.\n\n\n * Cookie\n__hssrc\n * Duration\nsession\n * Description\nThis cookie is set by Hubspot whenever it changes the session cookie. The __hssrc cookie set to 1 indicates that the user has restarted the browser, and if the cookie does not exist, it is assumed to be a new session.\n\n\n * Cookie\n__hssc\n * Duration\n1 hour\n * Description\nHubSpot sets this cookie to keep track of sessions and to determine if HubSpot should increment the session number and timestamps in the __hstc cookie.\n\n\n * Cookie\nwpEmojiSettingsSupports\n * Duration\nsession\n * Description\nWordPress sets this cookie when a user interacts with emojis on a WordPress site. It helps determine if the user's browser can display emojis properly.\n\n\n * Cookie\nBCSessionID\n * Duration\n1 year 1 month 4 days\n * Description\nBlueconic sets this cookie as a unique identifier for the BlueConic profile.\n\n\n * Cookie\n_octo\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nlogged_in\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\n__Secure-YEC\n * Duration\npast\n * Description\nYouTube sets this cookie to stores the user's video player preferences using embedded YouTube video\n\n\n * Cookie\n__eoi\n * Duration\n6 months\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTGCORS\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\nlogin-status-p\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTG\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\ncsrf_token\n * Duration\nsession\n * Description\nNo description available.\n\n\n * Cookie\ntoken_v2\n * Duration\n1 day\n * Description\nDescription is currently not available.\n\n\n * Cookie\nD\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nPHPSESSID\n * Duration\nsession\n * Description\nThis cookie is native to PHP applications. The cookie stores and identifies a user's unique session ID to manage user sessions on the website. The cookie is a session cookie and will be deleted when all the browser windows are closed.\n\n\n * Cookie\nVISITOR_PRIVACY_METADATA\n * Duration\n6 months\n * Description\nYouTube sets this cookie to store the user's cookie consent state for the current domain. \n\n\n * Cookie\ncookietest\n * Duration\nsession\n * Description\nThe cookietest cookie is typically used to determine whether the user's browser accepts cookies, essential for website functionality and user experience.\n\n\n * Cookie\n__Host-airtable-session\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\n__Host-airtable-session.sig\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\nm\n * Duration\n1 year 1 month 4 days\n * Description\nStripe sets this cookie for fraud prevention purposes. It identifies the device used to access the website, allowing the website to be formatted accordingly.\n\n\n * Cookie\nBIGipServer*\n * Duration\nsession\n * Description\nMarketo sets this cookie to collect information about the user's online activity and build a profile about their interests to provide advertisements relevant to the user. \n\n\n * Cookie\n__cfruid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to identify trusted web traffic.\n\n\n * Cookie\n_GRECAPTCHA\n * Duration\n6 months\n * Description\nGoogle Recaptcha service sets this cookie to identify bots to protect the website against malicious spam attacks.\n\n\n * Cookie\n__Secure-YNID\n * Duration\n6 months\n * Description\nGoogle cookie used to protect user security and prevent fraud, especially during the login process.\n\n\n * Cookie\ncookieyes-consent\n * Duration\n1 year\n * Description\nCookieYes sets this cookie to remember users' consent preferences so that their preferences are respected on subsequent visits to this site. It does not collect or store any personal information about the site visitors.\n\n\nFunctional\nFunctional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features.\n * Cookie\nlidc\n * Duration\n1 day\n * Description\nLinkedIn sets the lidc cookie to facilitate data center selection.\n\n\n * Cookie\nbrw\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nbrwConsent\n * Duration\n5 minutes\n * Description\nDescription is currently not available.\n\n\n * Cookie\nWMF-Uniq\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nloom_anon_comment\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nloom_referral_video\n * Duration\nsession\n * Description\nDescription is currently not available.\n\n\n * Cookie\nVISITOR_INFO1_LIVE\n * Duration\n6 months\n * Description\nA cookie set by YouTube to measure bandwidth that determines whether the user gets the new or old player interface.\n\n\n * Cookie\nyt-remote-connected-devices\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nytidb::LAST_RESULT_ENTRY_KEY\n * Duration\nNever Expires\n * Description\nThe cookie ytidb::LAST_RESULT_ENTRY_KEY is used by YouTube to store the last search result entry that was clicked by the user. This information is used to improve the user experience by providing more relevant search results in the future.\n\n\n * Cookie\nyt-remote-device-id\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-name\n * Duration\nsession\n * Description\nThe yt-remote-session-name cookie is used by YouTube to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\nyt-remote-fast-check-period\n * Duration\nsession\n * Description\nThe yt-remote-fast-check-period cookie is used by YouTube to store the user's video player preferences for embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-app\n * Duration\nsession\n * Description\nThe yt-remote-session-app cookie is used by YouTube to store user preferences and information about the interface of the embedded YouTube video player.\n\n\n * Cookie\nyt-remote-cast-available\n * Duration\nsession\n * Description\nThe yt-remote-cast-available cookie is used to store the user's preferences regarding whether casting is available on their YouTube video player.\n\n\n * Cookie\nyt-remote-cast-installed\n * Duration\nsession\n * Description\nThe yt-remote-cast-installed cookie is used to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\ncp_session\n * Duration\n3 months\n * Description\nCodepen sets this cookie for Help systems found in the website.\n\n\n * Cookie\nloid\n * Duration\n1 year 1 month 4 days\n * Description\nThis cookie is set by the Reddit. The cookie enables the sharing of content from the website onto the social media platform.\n\n\nAnalytics\nAnalytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc.\n * Cookie\n__hstc\n * Duration\n6 months\n * Description\nHubspot set this main cookie for tracking visitors. It contains the domain, initial timestamp (first visit), last timestamp (last visit), current timestamp (this visit), and session number (increments for each subsequent session).\n\n\n * Cookie\nhubspotutk\n * Duration\n6 months\n * Description\nHubSpot sets this cookie to keep track of the visitors to the website. This cookie is passed to HubSpot on form submission and used when deduplicating contacts.\n\n\n * Cookie\n_ga\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to calculate visitor, session and campaign data and track site usage for the site's analytics report. The cookie stores information anonymously and assigns a randomly generated number to recognise unique visitors.\n\n\n * Cookie\n_ga_*\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to store and count page views.\n\n\n * Cookie\n__Host-psifi.analyticsTrace\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n__Host-psifi.analyticsTraceV2\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n_gh_sess\n * Duration\nsession\n * Description\nGitHub sets this cookie for temporary application and framework state between pages like what step the user is on in a multiple step form.\n\n\n * Cookie\nYSC\n * Duration\nsession\n * Description\nYSC cookie is set by Youtube and is used to track the views of embedded videos on Youtube pages.\n\n\n * Cookie\najs_anonymous_id\n * Duration\n1 year\n * Description\nThis cookie is set by Segment to count the number of people who visit a certain site by tracking if they have visited before.\n\n\n * Cookie\nvuid\n * Duration\n1 year 1 month 4 days\n * Description\nVimeo installs this cookie to collect tracking information by setting a unique ID to embed videos on the website. \n\n\nPerformance\nPerformance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors.\n * Cookie\nAWSALB\n * Duration\n7 days\n * Description\nAWSALB is an application load balancer cookie set by Amazon Web Services to map the session to the target.\n\n\n * Cookie\nacq\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nacq.sig\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nptc\n * Duration\n2 years\n * Description\nNo description available.\n\n\nAdvertisement\nAdvertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns.\n * Cookie\nmuc_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to collect user behaviour and interaction data to optimize the website.\n\n\n * Cookie\nguest_id_marketing\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\nguest_id_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\npersonalization_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to integrate and share features for social media and also store information about how the user uses the website, for tracking and targeting.\n\n\n * Cookie\nguest_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor. It registers if a user is signed in to the Twitter platform and collects information about ad preferences.\n\n\n * Cookie\nbcookie\n * Duration\n1 year\n * Description\nLinkedIn sets this cookie from LinkedIn share buttons and ad tags to recognize browser IDs.\n\n\n * Cookie\n__Secure-ROLLOUT_TOKEN\n * Duration\n6 months\n * Description\nYouTube sets this cookie to manage feature rollout and experimentation. It helps Google control which new features or interface changes are shown to users as part of testing and staged rollouts, ensuring consistent experience for a given user during an experiment.\n\n\n * Cookie\nyt.innertube::nextId\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nyt.innertube::requests\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nsession_tracker\n * Duration\nsession\n * Description\nThis cookie is set by the Reddit. This cookie is used to identify trusted web traffic. It also helps in adverstising on the website.\n\n\n * Cookie\nedgebucket\n * Duration\nsession\n * Description\nReddit sets this cookie to save the information about a log-on Reddit user, for the purpose of advertisement recommendations and updating the content.\n\n\n * Cookie\ndid\n * Duration\n1 year\n * Description\nArbor sets this cookie to show targeted ads to site visitors.This cookie expires after 2 months or 1 year.\n\n\nUncategorised\nOther uncategorised cookies are those that are being analysed and have not been classified into a category as yet.\nNo cookies to display.\nReject All Save My Preferences Accept All\nToggle Search[](https://towardsdatascience.com/)Toggle Mobile Navigation\n[](https://towardsdatascience.com/)\nPublish AI, ML & data-science insights to a global community of data professionals.\nSign in[Submit an Article](https://write.insightmediagroup.io)\n[Latest](https://towardsdatascience.com/latest/)[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Newsletter](https://towardsdatascience.com/tag/the-variable/)[ShipAI](https://towardsdatascience.com/shipai/)\n* * *\n[Write For TDS](https://towardsdatascience.com/submissions/)\n[LinkedIn](https://www.linkedin.com/company/towards-data-science/)[X](https://x.com/TDataScience)Toggle Search\n[Data Engineering](https://towardsdatascience.com/category/data-science/data-engineering/)\n# Optimizing Data Transfer in Distributed AI/ML Training Workloads\nA deep dive on data transfer bottlenecks, their identification, and their resolution with the help of NVIDIA Nsight™ Systems – part 3\n[Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\nJanuary 23, 202613 min read\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/1-2-1.jpeg)\n[Photo](https://unsplash.com/photos/landscape-photography-of-river-with-trees-ssEQdOiKd8U) by Jachan DeVol on Unsplash\nThis is the third part of a series of posts on optimizing data transfer using NVIDIA Nsight™ Systems (nsys) profiler. [Part one](https://towardsdatascience.com/optimizing-data-transfer-in-ai-ml-workloads/) focused on CPU-to-GPU data copies, and [part two](https://towardsdatascience.com/optimizing-data-transfer-in-batched-ai-ml-inference-workloads/) on GPU-to-CPU copies. In this post, we turn our attention to data transfer between GPUs.\nNowadays, it is quite common for AI/ML training — particularly of large models — to be distributed across multiple GPUs. While there are many different schemes for performing such distribution, what they all have in common is their reliance on the constant transfer of data — such as gradients, weights, statistics, and/or metrics — between the GPUs, throughout training. As with the other types of data transfer we analyzed in our previous posts, here too, a poor implementation could easily lead to under-utilization of compute resources and the unjustified inflation of training costs. Optimizing GPU-to-GPU communication is an active area of research and innovation involving both hardware and software development.\nIn this post, we will focus on the most common form of distributed training — data-distributed training. In data-distributed training, identical copies of the ML model are maintained on each GPU. Each input batch is evenly distributed among the GPUs, each of which executes a training step to calculate the local gradients. The local gradients are then shared and averaged across the GPUs, resulting in an identical gradient update to each of the model copies. Using [NVIDIA Nsight™ Systems (nsys)](https://docs.nvidia.com/nsight-systems/index.html) profiler, we will analyze the effect of the GPU-to-GPU transfer of the gradients on the runtime performance of training a toy model and assess a few techniques for reducing its overhead.\nLearn this step by step with the interactive [Data Engineer](https://roadmap.sh/data-engineer) roadmap.\n#### Disclaimers\nThe code we will share in this post is intended for demonstrative purposes; please do not rely on its accuracy or optimality. Please do not interpret our mention of any tool, framework, library, service, or platform as an endorsement of its use.\nThanks to [Yitzhak Levi](https://www.linkedin.com/in/yitzhak-levi-49a217201/) for his contributions to this post.\n## Instance Selection for Distributed Training\nIn a previous post, [Instance Selection for Deep Learning](https://towardsdatascience.com/instance-selection-for-deep-learning-7463d774cff0/), we discussed the importance of choosing an instance type that is most suitable for your AI/ML workload and the potential impact of that choice on the success of your project. When choosing an instance type for a workload that includes a lot of GPU-to-GPU traffic, you will want to pay attention to how such communication is carried out, including: the instance topology, GPU interconnects, maximal throughput, and latency.\nIn this post, we will limit our discussion to distribution between GPUs on a single instance. We will experiment with two instance types: an [Amazon EC2 g6e.48xlarge](https://aws.amazon.com/ec2/instance-types/g6e/) with 8 [NVIDIA L40S](https://www.nvidia.com/en-eu/data-center/l40s/) GPUs and an [Amazon EC2 p4d.24xlarge](https://aws.amazon.com/ec2/instance-types/p4/) with 8 [NVIDIA A100](https://www.nvidia.com/en-eu/data-center/a100/) GPUs. Each will run an [AWS Deep Learning (Ubuntu 24.04) AMI](https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html) with [PyTorch (2.8)](https://pytorch.org/get-started/locally/), [nsys-cli profiler](https://developer.nvidia.com/nsight-systems/get-started) (version 2025.6.1), and the [NVIDIA Tools Extension](https://docs.nvidia.com/tools-extension/index.html) (NVTX) library.\nOne of the primary differences between the two instance types is how the GPUs are connected: On the [g6e.48xlarge](https://aws.amazon.com/ec2/instance-types/g6e/) communication between the GPUs is over [PCI Express (PCIe)](https://en.wikipedia.org/wiki/PCI_Express), whereas [p4d.24xlarge](https://aws.amazon.com/ec2/instance-types/p4/) includes [NVIDIA NVLink™](https://www.nvidia.com/en-eu/data-center/nvlink/) — dedicated hardware for enabling high-throughput GPU-to-GPU communication. Communication over the PCIe bus is significantly slower than NVLink. While this may be sufficient for workloads with a low communication-to-compute ratio, it could be a performance killer for workloads with high communication rates.\nTo discover the topology of the instance types, we run the following command:\nbash\n\n```\nnvidia-smi topo -m\n```\n\nOn the [g6e.48xlarge](https://aws.amazon.com/ec2/instance-types/g6e/), we get the following results:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/2-2-1024x343.webp)\nGPU Topology of [g6e.48xlarge](https://aws.amazon.com/ec2/instance-types/g6e/) (by Author)\nGPUs on the same NUMA node are connected by a “NODE” link and GPUs on different NUMA nodes by a “SYS” link. Both links traverse the PCIe as well as another interconnect; neither is a direct connection (a.k.a., a “PIX” link). We will see later on how this can impact throughput performance.\nOn the [p4d.24xlarge](https://aws.amazon.com/ec2/instance-types/p4/), every pair of GPUs is linked by a dedicated NVLink connection:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/3-2-1024x341.webp)\nGPU Topology of [p4d.24xlarge](https://aws.amazon.com/ec2/instance-types/p4/) (by Author)\n## A Toy Model\nTo facilitate our discussion, we build a toy data-distributed training experiment.\nWe choose a model with a relatively high communication-to-compute ratio — a [Vision Transformer (ViT)-L/32](https://docs.pytorch.org/vision/main/models/generated/torchvision.models.vit_l_32.html#torchvision.models.vit_l_32) image-classification model with roughly 306 million parameters — and a synthetic dataset that we will use to train it:\npython\n\n```\nimport os, time, torch, nvtx\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.distributed as dist\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision.models import vit_l_32\n\nWORLD_SIZE = int(os.environ.get(\"WORLD_SIZE\", 1))\nBATCH_SIZE = 32\nIMG_SIZE = 224\nWARMUP_STEPS = 10\nPROFILE_STEPS = 3\nCOOLDOWN_STEPS = 1\nTOTAL_STEPS = WARMUP_STEPS + PROFILE_STEPS + COOLDOWN_STEPS\nN_WORKERS = 8\n\ndef get_model():\n return vit_l_32(weights=None)\n\n# A synthetic dataset with random images and labels\nclass FakeDataset(Dataset):\n def __len__(self):\n return TOTAL_STEPS * BATCH_SIZE * WORLD_SIZE\n\n def __getitem__(self, index):\n img = torch.randn((3, IMG_SIZE, IMG_SIZE))\n label = torch.randint(0, 1000, (1,)).item()\n return img, label\n\ndef get_data_iter(rank, world_size):\n dataset = FakeDataset()\n\n sampler = DistributedSampler(dataset, num_replicas=world_size,\n rank=rank, shuffle=True)\n\n train_loader = DataLoader(dataset, batch_size=BATCH_SIZE,\n sampler=sampler, num_workers=N_WORKERS, \n pin_memory=True)\n\n return iter(train_loader)\n```\n\nWe define a utility function that we will use to set up PyTorch [DistributedDataParallel](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) (DDP) training. We configure the PyTorch [distributed communication package](https://docs.pytorch.org/docs/stable/distributed.html) to use the [NVIDIA Collective Communications Library](https://developer.nvidia.com/nccl) (NCCL) as its [communication backend](https://docs.pytorch.org/tutorials/intermediate/dist_tuto.html#communication-backends) and wrap the model in a [DistributedDataParallel](https://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) container.\npython\n\n```\ndef configure_ddp(model, rank):\n dist.init_process_group(\"nccl\")\n model = DDP(model,\n device_ids=[rank],\n bucket_cap_mb=2000)\n return model\n```\n\nNote that we naively configure the DDP bucket capacity to 2 GB. DDP groups the model’s gradients into buckets and performs gradient reduction of each bucket in a separate command. The implication of setting the bucket capacity to 2 GB is that all of the ~306 million (FP32) gradients will fit in a single bucket (306 million x 4 bytes per gradient = ~1.22 GB) and reduction will only occur once all gradients have been calculated.\nAs in our previous posts, we schedule nsys profiler [programmatically](https://github.com/pytorch/pytorch/blob/v2.9.0/torch/cuda/profiler.py#L50) and wrap the different portions of the training step with color-coded [NVTX](https://nvtx.readthedocs.io/en/latest/index.html) annotations:\npython\n\n```\ndef train(use_ddp=False):\n # detect the env vars set by torchrun\n rank = int(os.environ.get(\"RANK\", 0))\n local_rank = int(os.environ.get(\"LOCAL_RANK\", 0))\n torch.cuda.set_device(local_rank)\n\n model = get_model().to(local_rank)\n criterion = nn.CrossEntropyLoss().to(local_rank)\n\n if use_ddp:\n model = configure_ddp(model, rank)\n\n optimizer = optim.SGD(model.parameters())\n\n data_iter = get_data_iter(rank, WORLD_SIZE)\n\n model.train()\n\n for i in range(TOTAL_STEPS):\n\n # Schedule Profiling\n if i == WARMUP_STEPS:\n torch.cuda.synchronize()\n start_time = time.perf_counter()\n torch.cuda.profiler.start()\n elif i == WARMUP_STEPS + PROFILE_STEPS:\n torch.cuda.synchronize()\n torch.cuda.profiler.stop()\n end_time = time.perf_counter()\n\n with nvtx.annotate(f\"Batch {i}\", color=\"blue\"):\n\n with nvtx.annotate(\"get batch\", color=\"red\"):\n data, target = next(data_iter)\n data = data.to(local_rank, non_blocking=True)\n target = target.to(local_rank, non_blocking=True)\n with nvtx.annotate(\"forward\", color=\"green\"):\n output = model(data)\n loss = criterion(output, target)\n\n with nvtx.annotate(\"backward\", color=\"purple\"):\n optimizer.zero_grad()\n loss.backward()\n\n with nvtx.annotate(\"optimizer step\", color=\"yellow\"):\n optimizer.step()\n\n if use_ddp:\n dist.destroy_process_group()\n\n if rank == 0:\n total_time = end_time - start_time\n print(f\"Throughput: {PROFILE_STEPS/total_time:.2f} steps/sec\")\n\n\nif __name__ == \"__main__\":\n # enable ddp if run with torchrun\n train(use_ddp=\"RANK\" in os.environ)\n```\n\nWe set the `NCCL_DEBUG` environment variable for visibility into how NCCL sets up the transport links between the GPUs.\nbash\n\n```\nexport NCCL_DEBUG=INFO\n```\n\n### Single-GPU Performance\nWe begin our experimentation by running our script on a single GPU without DDP:\nbash\n\n```\nnsys profile \\\n --capture-range=cudaProfilerApi \\\n --capture-range-end=stop \\\n --trace=cuda,nvtx,osrt,nccl \\\n --output=baseline \\\n python train.py\n```\n\nNote the inclusion of `nccl`in the _trace_ section; this will be critical for analyzing the GPU-to-GPU communication in the multi-GPU experiments.\nThe throughput of the baseline experiment is 8.91 steps per second on the L40S GPU and 5.17 steps per second on the A100: The newer NVIDIA Ada Lovelace Architecture performs better than the NVIDIA Ampere Architecture on our toy model. In the image below we show the timeline view for the L40S experiment.\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/4-2-1024x215.webp)\nSingle L40S GPU Nsight Systems Trace (by Author)\nIn this post we will focus on the CUDA portion of the timeline. Looking at the NVTX row, we see the recurring red (CPU-to-GPU copy), green (forward), purple (backward), and yellow (optimizer update) pattern that makes up our train step. Note, however, that the purple appears as just a tiny blip while, in practice, the backward pass fills the entire gap between the forward and optimizer blocks: The NVTX library appears to have captured only the initial launch of the autograd graph.\nWe will use this trace as a comparative baseline for our next experiments.\n### DDP with 1 GPU\nWe assess the impact of the DDP wrapper by running [torchrun](https://docs.pytorch.org/docs/stable/elastic/run.html) on a single GPU:\nbash\n\n```\ntorchrun --nproc_per_node=1 \\\n --no-python \\\n nsys profile \\\n --capture-range=cudaProfilerApi \\\n --capture-range-end=stop \\\n --trace=cuda,nvtx,nccl,osrt \\\n --output=ddp-1gpu \\\n python train.py\n```\n\nThe resulting throughput drops to 8.40 steps per second on the L40S GPU and 5.04 steps per second on the A100. Even in the absence of any cross-GPU communication, DDP introduces overhead that can decrease throughput by ~3–7%. The important lesson from this is to always ensure that single GPU training is run without DDP.\nThe nsys trace of the L40S experiment confirms the presence of the DDP overhead:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/5-2-1024x265.webp)\nSingle L40S GPU With DDP - Profiler Trace (by Author)\nThe main change from the training step in the previous trace is a large chunk of device-to-device (_DtoD_) memory copies at the end of the backward block and just before the optimizer block (highlighted in the trace above). This is the DDP in action: Even in the absence of a cross GPU gradient reduction, DDP prepares the local gradients in a dedicated memory block for reduction and then copies the results back into the grad property of each parameter in preparation for the gradient update. (Note that the _DtoD_ copies are between two memory locations on the same GPU — not between two different GPUs).\n### DDP With Multiple GPUs\nNext, we assess the impact of gradient sharing between 2, 4, and 8 GPUs:\nbash\n\n```\ntorchrun --nproc_per_node=8 \\\n --no-python \\\n nsys profile \\\n --capture-range=cudaProfilerApi \\\n --capture-range-end=stop \\\n --trace=cuda,nvtx,nccl,osrt \\\n --output=ddp-8gpu_%q{RANK} \\\n python train.py\n```\n\nThe table below captures impact on the training throughput on the [g6e.48xlarge](https://aws.amazon.com/ec2/instance-types/g6e/) and the [p4d.24xlarge](https://aws.amazon.com/ec2/instance-types/p4/):\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/6-3.webp)\nThroughput Performance of DDP Training (by Author)\nDDP training performance plummets on the g6e.48xlarge instance: The 8-GPU throughput is more than 6 times slower than the single-GPU result. The NCCL logs include multiple lines describing the communication paths between the GPUs, e.g.:\nbash\n\n```\nNCCL INFO Channel 00 : 0[0] -> 1[1] via SHM/direct/direct\n```\n\nThis indicates that the data transfer between each two GPUs passes through CPU shared memory which greatly limits the communication bandwidth.\nOn the p4d.24xlarge, in contrast, where the NVLink connections allow for direct peer-to-peer (P2P) communication, the 8-GPU throughput is just 8% lower than the baseline result:\nbash\n\n```\nNCCL INFO Channel 00/0 : 0[0] -> 1[1] via P2P/CUMEM/read\n```\n\nEven though each individual L40S outperforms the A100 on our toy model by 72%, the inclusion of NVLink makes the p4d.24xlarge more optimal than the g6e.48xlarge for running DDP over 8 GPUs.\nThe data transfer bottleneck on the g6e.48xlarge can be easily seen in the nsys trace. Here we use the “multiple view” option to display the activity of all the DDP processes in a single timeline:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/7-3-1024x311.webp)\n8-GPU DDP on g6e.48xlarge - Profiler Trace (by Author)\nThe gradient reduction occurs in the _ncclAllReduce_ call on the NCCL row. This comes just after the local gradient calculation has completed and just before the _DtoD_ memory operation discussed above that copies the reduced gradients back to the _grad_ field of each parameter. On the g6e.48xlarge the NCCL operation dominates the backward pass, accounting for approximately 84% of the overall step time (587 out of 701 milliseconds).\nOn the p4d.24xlarge, the NCCL call takes up a much smaller portion of the training step:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/8-2-1024x342.webp)\n8-GPU DDP on p4d.24xlarge - Profiler Trace (by Author)\nIn the next sections we will discuss a few DDP optimization techniques and assess their impact on runtime performance. We will limit our scope to PyTorch-level optimizations, leaving other methods (e.g., NCCL/OS/network tuning) for another post.\nA general approach for dealing with communication bottlenecks is to reduce the frequency of communication. In DDP workloads this can be achieved by increasing the data batch size (i.e., increasing the compute per step) or, if the memory capacity forbids this, apply gradient accumulation — instead of applying the gradient reduction and update every step, accumulate the local gradients for _N_ steps and apply the reduction every _N_ th step. Both techniques increase the _effective_ batch size — the number of overall samples between gradient updates which can impact model convergence and may require tuning of the optimization parameters. In our discussion we assume the _effective_ batch size is fixed.\nWe will cover four techniques. The first two focus on reducing DDP overhead. The final two directly address the data transfer bottleneck.\n## Optimization 1: Static Graph Declaration\nOur first change is to pass `static_graph=True` to the [DDP](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) container. Generally speaking, models might include parameters for which gradients are not calculated at every step — referred to as “unused parameters” in the DDP documentation. This is common in models with conditional logic. DDP includes dedicated mechanisms for identifying and handling unused parameters. In the case of our toy model, all of the gradients are calculated in each step — our graph is “static”. Explicitly declaring our graph as static reduces the overhead associated with handling a dynamic gradient set.\npython\n\n```\ndef configure_ddp(model, rank):\n dist.init_process_group(\"nccl\")\n model = DDP(model,\n device_ids=[rank],\n static_graph=True,\n bucket_cap_mb=2000)\n return model\n```\n\nIn the case of our toy model, the impact of this change is negligible on both the g6e.48xlarge and the p4d.24xlarge. Without any further delay, we proceed to the next optimization.\n## Optimization 2: Increase Memory Efficiency\nOur second technique addresses the large chunk of _DtoD_ memory copying we identified above. Instead of copying the gradients to and from the NCCL communication memory blocks, we can explicitly set the parameter gradients to point directly to the NCCL communication buffers. Consequently, the same memory is used to store the local gradients, to perform the gradient reduction, and to apply the gradient updates. This is configured by passing `gradient_as_bucket_view=True` to the [DDP](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) container:\npython\n\n```\ndef configure_ddp(model, rank):\n dist.init_process_group(\"nccl\")\n model = DDP(model,\n device_ids=[rank],\n static_graph=True,\n gradient_as_bucket_view=True,\n bucket_cap_mb=2000)\n return model\n```\n\nIn the trace below, captured on the p4d.24xlarge, we no longer see the block of _DtoD_ memory copy between the all-reduce and (yellow) optimizer steps:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/9-1-1024x348.webp)\nMemory-Optimized DDP on p4d.24xlarge - Profiler Trace (by Author)\nIn the case of our toy example, this optimization boosts performance by a modest 1%.\n## Optimization 3: Gradient Compression\nA common approach for addressing communication bottlenecks is to apply compression algorithms to reduce the size of the payload. PyTorch DDP provides a number of [dedicated communication hooks](https://docs.pytorch.org/docs/stable/ddp_comm_hooks.html) that automate compression of gradients before NCCL reduction and decompression afterward. In the code block below, we apply [bfloat16 gradient compression](https://docs.pytorch.org/docs/stable/ddp_comm_hooks.html#torch.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_hook):\npython\n\n```\nimport torch.distributed.algorithms.ddp_comm_hooks.default_hooks as ddp_hks\n\ndef configure_ddp(model, rank):\n dist.init_process_group(\"nccl\")\n model = DDP(model,\n device_ids=[rank],\n static_graph=True,\n gradient_as_bucket_view=True,\n bucket_cap_mb=2000)\n\n model.register_comm_hook(state=None, hook=ddp_hks.bf16_compress_hook)\n return model\n```\n\nOn the heavily bottlenecked g6e.48xlarge instance bfloat16 compression results in a considerable 65% speed-up! The nsys trace shows the reduced-sized NCCL call as well as the newly introduced compression operations:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/10-1-1024x341.webp)\nBF16-Compressed DDP on g6e.48xlarge - Profiler Trace (by Author)\nOn the p4d.24xlarge the overhead of the compression operations outweigh the gains in the communication speed, leading to an overall reduction in throughput:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/11-1-1024x280.webp)\nBF16-Compressed DDP on p4d.24xlarge - Profiler Trace (by Author)\nPyTorch offers a more aggressive compression algorithm than bfloat16, called [PowerSGD](https://docs.pytorch.org/docs/stable/ddp_comm_hooks.html#powersgd-communication-hook). Below we present a naive usage — in practice this can require a lot of tuning. Please see the documentation for details:\npython\n\n```\nfrom torch.distributed.algorithms.ddp_comm_hooks import powerSGD_hook \n\ndef configure_ddp(model, rank):\n dist.init_process_group(\"nccl\")\n model = DDP(model,\n device_ids=[rank],\n static_graph=True,\n gradient_as_bucket_view=True,\n bucket_cap_mb=2000)\n\n state = powerSGD_hook.PowerSGDState(\n process_group=None\n )\n\n model.register_comm_hook(\n state,\n hook=powerSGD_hook.powerSGD_hook\n )\n return model\n```\n\n[PowerSGD](https://docs.pytorch.org/docs/stable/ddp_comm_hooks.html#powersgd-communication-hook) has a dramatic impact on the g6e.48xlarge instance, increasing throughput all the way back up to 7.5 steps per second — more than five times faster than the baseline result! Note the reduction in NCCL kernel size in the resultant trace:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/12-1024x424.webp)\nPowerSGD-Compressed DDP on g6e.48xlarge - Profiler Trace (by Author)\nIt is important to note that these compression algorithms are precision-lossy and should be used with caution, as they may impact your model convergence.\n## Optimization 4: Parallelize Gradient Reduction\nDDP groups parameters into multiple buckets and triggers gradient reduction of each bucket independently — as soon as the bucket is filled. This allows for running gradient reduction of filled buckets while gradient calculation (of other buckets) is still ongoing. The degree of overlap depends on the number of DDP buckets which we control via the _bucket_cap_mb_ setting of the DDP container. Recall that in our initial implementation, we explicitly set this to _2000_(2 GB) which (given the size of our model) translated to a single DDP bucket. The table below shows the throughput for different values of _bucket_cap_mb_. The optimal setting will vary based on the details of the model and runtime environment.\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/13-2.webp)\nDDP Throughput vs. Bucket Capacity (by Author)\nNote the significant 12% improvement when using multiple buckets on the g6e.48xlarge with BF16 compression. PowerSGD, on the other hand, works best when applied to a single bucket.\nIn the image below, captured on a p4d.24xlarge with _bucket_cap_mb_ set to _100_ , we can see the impact of this optimization on the profiler trace:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/14-1-1024x350.webp)\nDDP With Multiple Buckets on p4d.24xlarge - Profiler Trace (by Author)\nIn place of the single NCCL all-reduce operation, we now have 11 smaller blocks running in parallel with the local gradient computation.\n## Results\nWe summarize the results of our experiments in the following table:\n![](https://assets.insightmediagroup.io/media/wp-content/uploads/2026/01/15.webp)\nExperiment Results (by Author)\nOn the p4d.24xlarge, where the data is transferred over NVLink, the overall impact was a modest 4% speed-up. But on the g6e.48xlarge the gains were significant and mostly due to gradient compression — and 86% boost for the BF16 scheme and an over 5X improvement for (the naive implementation of) PowerSGD.\nImportantly, these results can vary considerably based on the model architecture and runtime environment.\n## Summary\nThis concludes our three-part series on identifying and solving data transfer bottlenecks with NVIDIA Nsight™ Systems (nsys) profiler. In each of the posts we demonstrated the tendency of data transfers to introduce performance bottlenecks and resource under-utilization. In each case, we used nsys profiler to identify the bottlenecks and measure the impact of different optimization techniques.\nThe accelerations we achieved in each of the use-cases that we studied strengthen the importance of integrating the regular use of tools such as nsys profiler into the AI/ML development workflow and highlight the opportunity — even for non-CUDA specialists — to achieve meaningful performance gains and AI/ML cost reductions.\n## Register Now to Continue Reading\n### Unlock Towards Data Science with a Free Account\nAll articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more.\nRegister for freeSign in to your account\n* * *\n![](https://assets.insightmediagroup.io/media/1787657094227_7ymyo0.jpeg)\nWritten By\n## [Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\n[See all from Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\n[Deep Dives](https://towardsdatascience.com/tag/deep-dives/)[Artificial Intelligence](https://towardsdatascience.com/tag/artificial-intelligence/)[Deep Learning](https://towardsdatascience.com/tag/deep-learning/)[Nvidia Gpu](https://towardsdatascience.com/tag/nvidia-gpu/)[Pytorch](https://towardsdatascience.com/tag/pytorch/)\nShare This Article\n * [](https://www.facebook.com/sharer/sharer.php?u=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-data-transfer-in-distributed-ai-ml-training-workloads%2F&title=Optimizing%20Data%20Transfer%20in%20Distributed%20AI%2FML%20Training%20Workloads)\n * [](https://www.linkedin.com/shareArticle?mini=true&url=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-data-transfer-in-distributed-ai-ml-training-workloads%2F&title=Optimizing%20Data%20Transfer%20in%20Distributed%20AI%2FML%20Training%20Workloads)\n * [](https://x.com/share?url=https%3A%2F%2Ftowardsdatascience.com%2Foptimizing-data-transfer-in-distributed-ai-ml-training-workloads%2F&text=Optimizing%20Data%20Transfer%20in%20Distributed%20AI%2FML%20Training%20Workloads)\n\n\nTowards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program.\n[Write for TDS](https://towardsdatascience.com/submissions/)\n## Related Articles\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2021/04/1-bz4_az8b7qGWwvONptkuA-scaled.jpeg)](https://towardsdatascience.com/all-data-integrations-should-use-change-data-capture-a1d207091773/)\n## [All Data Integrations Should Use Change Data Capture](https://towardsdatascience.com/all-data-integrations-should-use-change-data-capture-a1d207091773/)\n[Data Engineering](https://towardsdatascience.com/category/data-science/data-engineering/)\nCDC elevates the game for data integrations. But it is hard to to get right - learn what it takes to build data integrations with CDC.\n[Raghotham Murthy](https://towardsdatascience.com/author/raghotham/)\nApril 21, 202111 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2022/11/1KhAwZR9hPYSKZrxLHNmRAQ-scaled.jpeg)](https://towardsdatascience.com/dbt-models-structure-c31c8977b5fc/)\n## [How to Structure Your dbt Projects and Data Models](https://towardsdatascience.com/dbt-models-structure-c31c8977b5fc/)\n[Programming](https://towardsdatascience.com/category/programming/)\nEnforcing meaningful structure to dbt data models\n[Giorgos Myrianthous](https://towardsdatascience.com/author/gmyrianthous/)\nNovember 9, 20227 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2021/09/1SUQPrn_C1xiyk-1tT5NSSQ-scaled.jpeg)](https://towardsdatascience.com/10-most-practical-data-science-skills-you-should-know-in-2022-9487d7750e8a/)\n## [10 Most Practical Data Science Skills You Should Know in 2022](https://towardsdatascience.com/10-most-practical-data-science-skills-you-should-know-in-2022-9487d7750e8a/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nSkills that will actually make you employable\n[Terence Shin, MSc, MBA](https://towardsdatascience.com/author/terenceshin/)\nSeptember 5, 20217 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/06/04AKD33eZ4Ry8XwU6.png)](https://towardsdatascience.com/how-to-identify-your-business-critical-data-c58506116a41/)\n## [How to identify your business-critical data](https://towardsdatascience.com/how-to-identify-your-business-critical-data-c58506116a41/)\n[Data Science](https://towardsdatascience.com/category/data-science/)\nPractical steps to identifying business-critical data models and dashboards and drive confidence in your data\n[Mikkel Dengsøe](https://towardsdatascience.com/author/mikldd/)\nJune 14, 20239 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2024/09/11-MRL4GnYOzWR8zE1oXjFA.jpeg)](https://towardsdatascience.com/to-mask-or-not-to-mask-the-effect-of-prompt-tokens-on-instruction-tuning-016f85fd67f4/)\n## [To Mask or Not to Mask: The Effect of Prompt Tokens on Instruction Tuning](https://towardsdatascience.com/to-mask-or-not-to-mask-the-effect-of-prompt-tokens-on-instruction-tuning-016f85fd67f4/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nImplementing prompt-loss-weight, and why we should replace prompt-masking with prompt-weighting\n[David Vaughn](https://towardsdatascience.com/author/davidsvaughn/)\nSeptember 30, 202431 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2022/04/0Th0kSqxgWXhWvIdT-scaled.jpg)](https://towardsdatascience.com/7-ways-to-improve-performance-of-your-data-warehouse-4544456bcc3/)\n## [7 Ways to Improve Performance of Your Data Warehouse](https://towardsdatascience.com/7-ways-to-improve-performance-of-your-data-warehouse-4544456bcc3/)\n[Data Science](https://towardsdatascience.com/category/data-science/)\nImprove Poor Performance of Your Data Warehouse by Applying These Steps\n[Hanzala Qureshi](https://towardsdatascience.com/author/hanzalaqureshi/)\nApril 26, 20226 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2022/12/0vgxI8urMr_dIBmlN-scaled.jpg)](https://towardsdatascience.com/easy-parquet-tutorial-best-practices-237955e46cb7/)\n## [Simple Parquet Tutorial and Best Practices](https://towardsdatascience.com/easy-parquet-tutorial-best-practices-237955e46cb7/)\n[Data Science](https://towardsdatascience.com/category/data-science/)\nHands-on tutorial for starting your Parquet learning\n[Arli](https://towardsdatascience.com/author/arli94/)\nDecember 16, 20228 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/01/1sv67yR1CN18Jbv5Ht5G4Qg-scaled.jpeg)](https://towardsdatascience.com/unit-testing-pyspark-code-using-pytest-b5ab2fd54415/)\n## [Unit testing PySpark code using Pytest](https://towardsdatascience.com/unit-testing-pyspark-code-using-pytest-b5ab2fd54415/)\n[Data Science](https://towardsdatascience.com/category/data-science/)\nWhen it comes to writing unit-tests for PySpark pipelines, writing focussed, fast, isolated and concise tests can be a challenge.\n[Julian West](https://towardsdatascience.com/author/julianwest155/)\nJanuary 16, 202311 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2024/08/1j4Vw22q-llHYhh6YH2IH5w.png)](https://towardsdatascience.com/structured-state-space-models-visually-explained-86cfe2757386/)\n## [Structured State Space Models Visually Explained](https://towardsdatascience.com/structured-state-space-models-visually-explained-86cfe2757386/)\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\nPart 2 - Towards Mamba State Space Models for Images, Videos and Time Series\n[Sascha Kirch](https://towardsdatascience.com/author/saschakirch/)\nAugust 22, 202418 min read\n\n\n * [YouTube](https://www.youtube.com/c/TowardsDataScience)\n * [X](https://x.com/TDataScience)\n * [LinkedIn](https://www.linkedin.com/company/towards-data-science/)\n * [Threads](https://www.threads.net/@towardsdatascience)\n * [Bluesky](https://bsky.app/profile/towardsdatascience.com)\n\n\n[](https://towardsdatascience.com/)\nYour home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals.\nSubscribe to Our Newsletter[Write For TDS](https://towardsdatascience.com/submissions/)•[About](https://towardsdatascience.com/about-towards-data-science-d691af11cc2f/)•[Advertise](https://contact.towardsdatascience.com/advertise-with-towards-data-science)•[Privacy Policy](https://towardsdatascience.com/privacy-policy/)•[Terms of Use](https://towardsdatascience.com/website-terms-of-use/)\n© Insight Media Group, LLC 2026•[The New Stack](https://thenewstack.io)•[roadmap.sh](https://roadmap.sh)•[Creator Network](https://insightmediagroup.io/creator-network)\n#### Recommended Articles \nClose\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500) Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1 ](https://towardsdatascience.com/javascript-fatigue-you-dont-need-js-to-build-chatgpt/ \"Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500) The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs ](https://towardsdatascience.com/subset-sum-problem-solved-in-linear-time-for-dense-enough-inputs/ \"The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500) Does Your Company Have a Data Strategy? ](https://towardsdatascience.com/does-your-company-have-a-data-strategy-e8944dbec4ae/ \"Does Your Company Have a Data Strategy?\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500) Productivity Tips, Data Career Insights, and Other Recent Must-Reads ](https://towardsdatascience.com/productivity-tips-data-career-insights-and-other-recent-must-reads-5c47eb7a2fe0/ \"Productivity Tips, Data Career Insights, and Other Recent Must-Reads\")\n\n\n", "body_base_path": "attachments/f29ca4fd35ae/pages/", "body_markdown_path": "attachments/f29ca4fd35ae/pages/archived-body.md", "local_html": "attachments/f29ca4fd35ae/pages/archived-body.html", "body_source": "prior_capture" }