DeepSeek: TorchFeather

Notes while following Umar Jamil's Building a distributed training framework from first principles and working through TorchFeather. I use these notes to explain the implementation to myself as I study it.

Output Initialization: Why 1/√dim?🔗

In DeepSeekV3Model.init_weights:

final_out_std = self.model_args.dim**-0.5

This sets the initial standard deviation of the output layer's weights. The output layer maps \( d = \texttt{dim} \) hidden features to vocabulary logits; final_out_std is not the measured standard deviation of those logits.

Goal: Preserve the Input Variance🔗

For one output logit, write \( y = \sum_{j=1}^{d} w_j x_j \), summing over \( d \) input features. We choose the weight scale so \( \operatorname{Var}(y) \approx \operatorname{Var}(x_j) \): one output should have roughly the variance of one input feature. This is the forward variance-preservation argument behind fan-in initialization; see Glorot and Bengio, Section 4.2.

Which Variables Are Independent?🔗

For the simple derivation, assume each input feature is centered with variance \( v \), and the initialized weights are independent of the input and of each other, with mean zero and variance \( \sigma_w^2 \).

x and y are dependent: y is computed from x. Independence here refers to the randomly initialized weights and the input, not to the input and output.

For one product, independence lets us factor its second moment, while zero means remove the mean term:

\operatorname{Var}(w_j x_j) = \mathbb{E}[w_j^2]\,\mathbb{E}[x_j^2] = \sigma_w^2 v

For a sum, variance generally includes covariance terms:

\begin{aligned} \operatorname{Var}(y) &= \sum_{j=1}^{d} \operatorname{Var}(w_j x_j) \\ &\quad + 2\sum_{j<k} \operatorname{Cov}(w_j x_j,\, w_k x_k) \end{aligned}

The independent, zero-mean weights make those cross-covariances zero at initialization. The input features themselves need not be independent for this cancellation.

Solve for the Weight Standard Deviation🔗

There are \( d \) products, so:

\operatorname{Var}(y) = d\,\sigma_w^2 v

To preserve a nonzero input variance \( v \):

\begin{aligned} d\,\sigma_w^2 v &= v \\ \sigma_w^2 &= \frac{1}{d} \\ \sigma_w &= \frac{1}{\sqrt{d}} \end{aligned}

The variance is \( 1/d \); the standard deviation passed to the initializer is its square root. For \( d = 4096 \), this gives \( \frac{1}{64} = 0.015625 \). Here \( d \) is the hidden width, not the vocabulary size.

Connect the Derivation to RMSNorm🔗

TorchFeather applies RMSNorm before the output projection. RMSNorm controls the mean square, without forcing the input mean to zero. Holding that hidden vector \( x \) fixed and varying the initial weights gives:

\operatorname{Var}_{W}(y\mid x) = \sigma_w^2\sum_{j=1}^{d} x_j^2 \approx \sigma_w^2 d

Thus \( \sigma_w^2 = 1/d \) keeps initial logit variance around 1 even without the centered-input simplification above. The code truncates the normal distribution at \( \pm 3\sigma \), so this is approximate scale control at initialization, not an exact variance constraint enforced throughout training.