THE EMBEDDING TABLE AUDIT

Difficulty: Medium Tags: Deep Learning, Embeddings, Memory Accounting

Story

A language model team estimates the memory cost of its token embedding table but forgets that an optional output embedding may be untied from the input embedding. Compute the parameter count and storage precisely.

The Math

An embedding table has vocabulary size V and dimension D. If input and output embeddings are tied, the table has V*D parameters; if untied, there are 2*V*D parameters. With b bytes per parameter:

M = N_params * b

Optimizer states, alignment, and other model weights are ignored.

Input Format

v, d, b (ints), tied (bool: True if tied, False if untied).

Output Format

(n_params, storage_bytes) as a tuple of ints.

Constraints

• 1 <= V, D <= 10^6, b in {1,2,4,8}.
• Exact 64-bit integer arithmetic.
• Time limit: 1.0 second.

Example

Input

V=50000, D=4096, b=2, tied=True

Output

(204800000, 409600000)

Explanation: Tied embeddings use 50000*4096=204800000 parameters; at 2 bytes each this is 409600000 bytes.