Memory Footprint Estimation
The problem, from first principles
02-parameter-counting answers "how many numbers does this model have." That's not yet "will this model fit on my GPU" — the actual memory cost depends on how many bytes each number takes, which depends on precision (float32 vs float16 vs int8), and on the fact that training needs far more than just the parameters themselves: gradients and optimizer state add up fast.
From theory to code
Implement estimate_memory_bytes(params, dtype="float32"), the raw storage cost of a parameter list at a given precision, and estimate_optimizer_memory_bytes(params, dtype="float32", optimizer="sgd"), the full training-time memory cost including gradients and optimizer state.
Constraints
dtypeis one of"float32"(4 bytes/element),"float16"/"bfloat16"(2 bytes/element),"int8"(1 byte/element).estimate_memory_bytesreturns total elements times bytes-per-element for the given dtype.estimate_optimizer_memory_bytesreturns parameter bytes + gradient bytes (always equal to parameter bytes) + optimizer state bytes, where optimizer state is0xparam bytes for"sgd",1xfor"momentum",2xfor"adam".- An empty
paramslist returns0bytes in both functions.
Hints
Open one at a time. Each gives away a little more than the last.
Hint 1
Reuse the same total-element-count idea as 02-parameter-counting, then multiply by however many bytes one element of dtype takes.
Hint 2
Every optimizer needs to store a gradient the same size as the parameters it's updating — that's always 1x extra, regardless of which optimizer. Beyond that, sgd needs nothing more, momentum needs one more same-sized buffer, adam needs two (its first and second moment estimates).
Keywords
No results yet
Run your code to check it against the examples, or submit to run every test.