2c:null 2d:["$","title",null,{"children":"PowerSGD Benchmark | Thijs Vogels | Observable"}] 32:T49b,md` # PowerSGD Benchmark I compared PyTorch's new [implementation of PowerSGD](https://github.com/pytorch/pytorch/blob/master/torch/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.py) against the original research code from the [paper](https://arxiv.org/abs/1905.13727). This new implementation based on [DistributedDataParallel](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html) makes it easier to use PowerSGD in existing applications: --- \`\`\`python import torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook as powerSGD state = PowerSGD.PowerSGDState( process_group=process_group, matrix_approximation_rank=4 ) ddp_model.register_comm_hook(state, PowerSGD.powerSGD_hook) \`\`\` --- The implementations behave similarly in terms of both speed and accuracy on the 16-worker Cifar-10 experiments from the paper. The experiments were executed on Google Cloud with \`n1-standard-4\` instances with one Nivdia K80 GPU per node. I used OpenMPI 4.0.3 with PMIx 3.1.5 and UCX as a communication backend in PyTorch 1.8.0a0+3397919. Each experiment was repeated twice with different seeds to gauge variance between runs. `33:T552,render(htm`

${["PyTorch DDP", "Paper"].map( implementation => htm`

${implementation}

${(implementation === "PyTorch DDP" ? [0, "FP16", 1, 2, 4] : [0, 1, 2, 4] ).map( rank => htm`

${ rank == 0 ? "Full precision" : rank === "FP16" ? "Float16" : `Rank ${rank}` }

${mean( batchTimes(implementation, rank, 16) ).toFixed(3)}s

` )}

`)34:T4eb,timingResults = FileAttachment("timingResults.json").json() // { // const data = await graphql(`{ // jobs(project: "powersgd", status: FINISHED) { // experiment // logs // config { key value } // } // }`); // const splitLines = str => str.split(/\r?\n/); // const filter = ({ experiment, logs, config }) => { // const configMap = Object.fromEntries(config.map(c => [c.key, c.value])); // const logLines = splitLines(logs); // const batchTimeLines = logLines.filter( // line => line.indexOf("- batch ") > -1 // ); // const batchTime = parseFloat( // batchTimeLines[batchTimeLines.length - 1].match(/([0-9.]+)s/)[1] // ); // return { // implementation: { // pytorch_varying_rank_ours: "Paper", // pytorch_varying_rank: "PyTorch DDP" // }[experiment], // batchTime, // n_workers: configMap["n_workers"], // rank: configMap["fp16_compression"] // ? "FP16" // : configMap["optimizer_reducer_rank"] || 0 // }; // }; // return data["jobs"] // .filter( // j => // ["pytorch_varying_rank_ours", "pytorch_varying_rank"].indexOf( // j.experiment // ) >= 0 // ) // .map(filter); // }35:T638,async function showconfig(df, blacklist = new Set(["distributed_init_file"])) { const ranges = new Map(); for (let entry of df) { const { config, job, time, value, epoch, steps, curve, ...otherProps } = entry; for (let [key, value] of Object.entries(job)) { if (["config", "timeseries", "id"].includes(key)) continue; if (!ranges.has(key)) { ranges.set(key, new Set()); } ranges.get(key).add(value); } for (let [key, value] of Object.entries(otherProps)) { if (!ranges.has(key)) { ranges.set(key, new Set()); } ranges.get(key).add(value); } for (let [key, value] of Object.entries(config)) { if (!ranges.has(key)) { ranges.set(key, new Set()); } ranges.get(key).add(value); } } const truncateText = (text, maxLength) => { if (text.length <= maxLength) { return text; } else { return text.slice(0, maxLength) + " ..."; } }; return render(htm`

${Array.from(ranges.entries()) .filter(([key]) => !blacklist.has(key)) .map( ([key, values]) => htm`` )}

${key.replace(/_/g, " ")}

${truncateText( Array.from(values) .map(v => (v == null ? "null" : v)) .sort(d3.ascending) .join(", "), 75 )}

`); }36:[{"id":"0dcbc9def41b8bd80c5ac3c50ab1aebce0c778b4e2e447cd25ba1c9ec635dc9514410215d6b93cbf1ff5eff8cc4e81c224ddce7ad5337148aa87f86095853d06","href":"https://static.observableusercontent.com/files/0dcbc9def41b8bd80c5ac3c50ab1aebce0c778b4e2e447cd25ba1c9ec635dc9514410215d6b93cbf1ff5eff8cc4e81c224ddce7ad5337148aa87f86095853d06","name":"results@3.json","type":"application/json","lastModified":1631774201235.094,"size":2447181},{"id":"6357ad844904c0983ed8fc9ae3ae70a91ff8a38174daa0fde3f36f37143b7e517883fdae6c3c9db6fb4649fc4031f2eab881b2100b2ca414edfc321cdc5579d8","href":"https://static.observableusercontent.com/files/6357ad844904c0983ed8fc9ae3ae70a91ff8a38174daa0fde3f36f37143b7e517883fdae6c3c9db6fb4649fc4031f2eab881b2100b2ca414edfc321cdc5579d8","name":"timingResults.json","type":"application/json","lastModified":1657138790145.76,"size":1902}] 30:{"id":"816c5b7aa109f6db","slug":"powersgd-benchmark","title":"PowerSGD Benchmark","thumbnail":"968ddd99294bac6345a9aad324ef8cd83a551ae9c2dd7db45b8e5e1c22d01075","access_level":"public","create_time":"$D2021-02-01T08:15:57.621Z","update_time":"$D2022-07-06T20:20:57.643Z","view_time":null,"first_public_time":"$D2021-02-01T09:50:47.684Z","has_unpublished_edits":null,"fork":false,"stars":0,"starred":false,"owner":{"type":"team","id":"63d4ae9351b957f2","login":"tvogels","name":"Thijs Vogels","avatar_url":"https://avatars-next.observableusercontent.com/caecb0802edbb2531d0d2883357c53d701439f3641bbc772df9d728d289a6905","tier":"starter_2024"},"version":467,"publish_version":467,"role":null,"authored":false,"authors":[{"type":"individual","id":"7ebda7165bbc067e","login":"tvogels","name":"Thijs Vogels","avatar_url":"https://avatars-next.observableusercontent.com/caecb0802edbb2531d0d2883357c53d701439f3641bbc772df9d728d289a6905","tier":null}],"members":[{"type":"individual","id":"7ebda7165bbc067e","login":"tvogels","name":"Thijs Vogels","avatar_url":"https://avatars-next.observableusercontent.com/caecb0802edbb2531d0d2883357c53d701439f3641bbc772df9d728d289a6905","tier":null,"role":"editor"}],"comments":[],"custom_thumbnail":false,"trashed":false,"lapsed":null,"subscribed":null,"forks":0,"collection":null,"type":"notebook","body":{"titling":"auto","title":"PowerSGD Benchmark","cells":[{"id":21,"value":"$32","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":105,"value":"md`---\n## Cifar-10`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":111,"value":"data = {\n // const addImplementation = (implementation) => (e) => ({\n // ...e,\n // implementation\n // });\n // const pytorch = await dataframe({\n // experiment: \"pytorch_varying_rank\",\n // project: \"powersgd\",\n // measurement: \"last_accuracy\",\n // tags: \"split=test\"\n // });\n // const original = await dataframe({\n // experiment: \"pytorch_varying_rank_ours\",\n // project: \"powersgd\",\n // measurement: \"last_accuracy\",\n // tags: \"split=test\"\n // });\n // return [\n // ...pytorch.map(addImplementation(\"PyTorch DDP\")),\n // ...original.map(addImplementation(\"Paper\"))\n // ];\n const { data, jobs, configs } = await FileAttachment(\"results@3.json\").json();\n for (let entry of data) {\n entry.job = jobs[entry.jobId];\n entry.config = configs[entry.jobId];\n }\n return data;\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":149,"value":"md`#### Test accuracy with 16 workers`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":123,"value":"{\n return render(\n plotgrid(data.filter(d => d.config.n_workers === 16), {\n colFn: d => `${d.implementation}`,\n rowLabelFn: r => `Test accuracy`,\n // rowFn: d => d.split,\n xFn: d => d.epoch,\n xLabel: \"Epochs\",\n ymin: 0.66,\n ymax: .95,\n legendLabel: \"Matrix approximation rank:\",\n colorFn: d =>\n d.config.optimizer_reducer_rank > 0\n ? `${d.config.optimizer_reducer_rank}`\n : d.config.fp16_compression\n ? \"Float16\"\n : \"Full precision\",\n curveFn: d => null,\n hMargin: 30,\n aspectRatio: 2\n })\n );\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":172,"value":"md`#### Test accuracy with 16 workers (zoomed in, from epoch 100)`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":166,"value":"{\n return render(\n plotgrid(data.filter(d => d.config.n_workers === 16), {\n colFn: d => `${d.implementation}`,\n rowLabelFn: r => `Test accuracy`,\n // rowFn: d => d.split,\n xFn: d => d.epoch,\n xLabel: \"Epochs\",\n xmin: 130,\n ymin: 0.92,\n ymax: .95,\n legendLabel: \"Matrix approximation rank:\",\n colorFn: d =>\n d.config.optimizer_reducer_rank > 0\n ? `${d.config.optimizer_reducer_rank}`\n : d.config.fp16_compression\n ? \"Float16\"\n : \"Full precision\",\n curveFn: d => null,\n hMargin: 30,\n aspectRatio: 2\n })\n );\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":187,"value":"md`#### Accuracy as a function of communication budget`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":175,"value":"{\n return render(\n plotgrid(data.filter(d => d.config.n_workers === 16), {\n colFn: d => `${d.implementation}`,\n rowLabelFn: r => `Test accuracy`,\n // rowFn: d => d.split,\n xFn: d => d.bits,\n xLabel: \"Bits communicated\",\n xmax: 2e11,\n ymin: 0.66,\n ymax: .95,\n legendLabel: \"Matrix approximation rank:\",\n colorFn: d =>\n d.config.optimizer_reducer_rank > 0\n ? `${d.config.optimizer_reducer_rank}`\n : d.config.fp16_compression\n ? \"Float16\"\n : \"Full precision\",\n curveFn: d => null,\n hMargin: 30,\n aspectRatio: 2\n })\n );\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":193,"value":"md`### Time per batch (16 workers, batch size 128)\nAverage over 300 epochs and over two independent experiments with different seeds.\n`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":236,"value":"$33","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":196,"value":"$34","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":300,"value":"function batchTimes(implementation, rank, n_workers) {\n return timingResults\n .filter(\n r =>\n r.implementation === implementation &&\n r.rank === rank &&\n r.n_workers === n_workers\n )\n .map(r => r.batchTime);\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":328,"value":"md`---\n\n## Detailed experimental settings`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":331,"value":"showconfig(data)","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":0,"value":"md`---\n\n## Imports`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":4,"value":"import { dataframe, graphql } from \"@tvogels/jobmonitor@248\"","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":6,"value":"import { render, plotgrid, htm } from \"b679d071989a855a@1154\"","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":10,"value":"d3 = require(\"d3@5\")","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":13,"value":"md`---\n\n## Library`","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":306,"value":"function mean(arr) {\n let sum = 0;\n arr.forEach(v => (sum += v));\n return sum / arr.length;\n}","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"},{"id":18,"value":"$35","pinned":false,"hidden":false,"mode":"ojs","author":"$undefined","summary":"$undefined","providerMetadata":"$undefined","attachments":"$undefined","output":"$undefined","database":"$undefined"}],"files":"$W36","stdlib":"1"}} 31:"$30:body"