Rotary embeddings not working because of dynamic shapes

[faresobeid]

I've tried to implement the RetNet model proposed in this paper: https://arxiv.org/pdf/2307.08621v4.pdf. Everything has worked well so far except for the XPOS part (or rotary embedding). In my code and other implementations, for the recurrent part of the model used for generating samples, the current timestep of the character has to be passed in and in the computation of the positional embeddings, an array based on that timestep is made.

class MSR_Recurrent(nn.Module):
    n_heads: int
    @nn.compact
    def __call__(self,x,past_state,i):
        batch_size, seq_len, d_model = x.shape
        gamma = 1 - jnp.exp((-5 - jnp.arange(0, self.n_heads)) * jnp.log(2))[:,None,None]
        q,k = jnp.dsplit(nn.Dense(2*d_model)(x),2)
        v,g = jnp.dsplit(nn.Dense(4*d_model)(x),2)
        xpos = XPOS(d_model)
        q = xpos(q,i+1)
        k = xpos(k,i+1, downscale=True)
        q = q.reshape((batch_size,seq_len,self.n_heads,-1)).transpose((0,2,1,3))
        k_T = k.reshape((batch_size,seq_len,self.n_heads,-1)).transpose((0,2,3,1))
        v = v.reshape((batch_size,seq_len,self.n_heads,-1)).transpose((0,2,1,3))
        current_state = gamma * past_state + jnp.matmul(k_T,v)
        out = jnp.matmul(q,current_state)
        concat_retention = (out.transpose(0, 2, 1, 3).reshape(batch_size * seq_len, -1)) 
        concat_retention = nn.GroupNorm(self.n_heads)(concat_retention).reshape(batch_size,seq_len,-1)
        out = nn.Dense(d_model)(nn.swish(g)*concat_retention)
        return out,current_state

Where i is the current timestep.

def fixed_pos_embedding(x):
    seq_len, dim = x.shape
    inv_freq = 1.0 / (10000 ** (jnp.arange(0, dim) / dim))
    sinusoid_inp = jnp.einsum("i , j -> i j", jnp.arange(0, seq_len, dtype=jnp.float32), inv_freq)
    return jnp.sin(sinusoid_inp), jnp.cos(sinusoid_inp)

def rotate_every_two(x):
    x1 = x[:, :, :, ::2]
    x2 = x[:, :, :, 1::2]
    x = jnp.stack((-x2, x1), axis=-1)
    return x.reshape(x.shape[:-2] + (-1,))  # in einsum notation: rearrange(x, '... d j -> ... (d j)')\

def duplicate_interleave(m):
    
    dim0 = m.shape[0]
    m = m.reshape((-1, 1))  # flatten the matrix
    m = m.repeat(1, 2)  # repeat all elements into the 2nd dimension
    m = m.reshape((dim0, -1))  # reshape into a matrix, interleaving the copy
    return m

def apply_rotary_pos_emb(x, sin, cos, scale=1):
    sin, cos = map(lambda t: duplicate_interleave(t * scale), (sin, cos))
    return (x * cos) + (rotate_every_two(x) * sin)


class XPOS(nn.Module):
    head_dim: int
    scale_base: int = 512
    
    @nn.compact
    def __call__(self, x, offset=0, downscale=False):
        length = x.shape[1]
        min_pos = 0
        max_pos = length + offset + min_pos
        scale = self.scale_base ** jnp.true_divide(jnp.arange(min_pos, max_pos, 1),(jnp.array(self.scale_base)))[:,None]
        
        sin, cos = fixed_pos_embedding(scale)
        if scale.shape[0] > length:
            scale = scale[-length:]
            sin = sin[-length:] 
            cos = cos[-length:]
        if downscale:
            scale = 1 / scale
        x = apply_rotary_pos_emb(x, sin, cos, scale)
        return x

And here is the sampling function:

def sample_step_scan(batch, model, max_len, params, past_states, key_sample):
    key = key_sample
    init_val = (batch,
               past_states,
               key)
    def loop_body(loop_carry,i):
        x,past_states,key = loop_carry
        key1,key2 = random.split(key)
        pred,past_state = model.apply_fn({'params': params},x,past_states,i)
        pred = pred[:,-1,:]
        probs = jax.nn.softmax(pred,axis=-1)
        pred_next = random.choice(key1,jnp.arange(0,65),shape=(1,1),p=probs[0])
        return (pred_next,past_state,key2),pred_next
    return jax.lax.scan(loop_body,init_val,jnp.arange(0,max_len))

This works fine without any errors by using a simple for loop for sampling, however it is extremely slow. However, when switching over to jax.lax.scan() or jax.lax.fori_loop() I get this error:

---------------------------------------------------------------------------
ConcretizationTypeError                   Traceback (most recent call last)
Cell In[34], line 1
----> 1 train(d_model=256,blocks=4,n_heads=4,vocab_size=65,batch_size=64,seq_len=256,seed=5)

Cell In[17], line 50, in train(d_model, blocks, n_heads, vocab_size, batch_size, seq_len, seed, max_len, eval_iters, max_iters)
     48 key_inference,subkey = random.split(key_inference)
     49 key_samples = random.split(subkey,jax.local_device_count())
---> 50 _,x = parallel_sample_step(batch, state_sample, max_len, params_sample, past_states, key_samples)
     51 print('------------------------------')
     52 print(decode(x[0][:,0,0].tolist()))

    [... skipping hidden 12 frame]

Cell In[12], line 52, in sample_step_scan(batch, model, max_len, params, past_states, key_sample)
     50     pred_next = random.choice(key1,jnp.arange(0,65),shape=(1,1),p=probs[0])
     51     return (pred_next,past_state,key2),pred_next
---> 52 return jax.lax.scan(loop_body,init_val,jnp.arange(0,max_len))

    [... skipping hidden 9 frame]

Cell In[12], line 47, in sample_step_scan.<locals>.loop_body(loop_carry, i)
     45 x,past_states,key = loop_carry
     46 key1,key2 = random.split(key)
---> 47 pred,past_state = model.apply_fn({'params': params},x,past_states,i)
     48 pred = pred[:,-1,:]
     49 probs = jax.nn.softmax(pred,axis=-1)

    [... skipping hidden 6 frame]

Cell In[8], line 16, in Retnet_Recurrent.__call__(self, x, past_states, i)
     11 x = nn.Embed(self.vocab_size,self.d_model)(x)
     14 for j in range(self.blocks):
---> 16     retention,state = MSR_Recurrent(self.n_heads)(nn.LayerNorm(use_bias=False)(x),past_states[:,j,:,:,:],i)
     17     x = retention + x
     18     x = FeedForward()(nn.LayerNorm(use_bias=False)(x)) + x

    [... skipping hidden 2 frame]

Cell In[6], line 11, in MSR_Recurrent.__call__(self, x, past_state, i)
      9 v,g = jnp.dsplit(nn.Dense(4*d_model)(x),2)
     10 xpos = XPOS(d_model)
---> 11 q = xpos(q,i+1)
     12 k = xpos(k,i+1, downscale=True)
     13 q = q.reshape((batch_size,seq_len,self.n_heads,-1)).transpose((0,2,1,3))

    [... skipping hidden 2 frame]

Cell In[33], line 38, in XPOS.__call__(self, x, offset, downscale)
     36 min_pos = 0
     37 max_pos = length + offset + min_pos
---> 38 scale = self.scale ** jnp.true_divide(jnp.arange(min_pos, max_pos, 1),self.scale_base)[:,None]
     40 sin, cos = fixed_pos_embedding(scale)
     42 if scale.shape[0] > length:

File /opt/conda/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:2292, in arange(start, stop, step, dtype)
   2290     start = core.concrete_or_error(None, start, "It arose in the jnp.arange argument 'start'")
   2291 util.check_arraylike_or_none("arange", None, stop, step)
-> 2292 stop = core.concrete_or_error(None, stop, "It arose in the jnp.arange argument 'stop'")
   2293 step = core.concrete_or_error(None, step, "It arose in the jnp.arange argument 'step'")
   2294 start_name = "stop" if stop is None and step is None else "start"

File /opt/conda/lib/python3.10/site-packages/jax/_src/core.py:1381, in concrete_or_error(force, val, context)
   1379     return force(val.aval.val)
   1380   else:
-> 1381     raise ConcretizationTypeError(val, context)
   1382 else:
   1383   return force(val)

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[].
It arose in the jnp.arange argument 'stop'
The error occurred while tracing the function loop_body at /tmp/ipykernel_28/3429937701.py:44 for scan. This concrete value was not available in Python because it depends on the value of the argument i.

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

I tried my way around this by manipulating the scan function but nothing has worked. Is there any way to do this without harming performance drastically. My guess is to not include the XPOS in the jit because it includes dynamic shapes. Surely theres a way to just include a 'jnp.arange(0,seq_len)' in my computation where seq_len changes.

[cgarciae]

Hey @faresobeid, take a look at this implementation of rotary embeddings:

https://github.com/google/flax/blob/nnx/flax/experimental/nnx/examples/07_transformer.py#L131-L157

[faresobeid]

Thank you! The error is not appearing anymore, however my samples are now much worse than before although the loss has decreased. This is what my training looks like after using these rotary embeddings:

step 0: train loss 5.5700, val loss 5.5890
------------------------------
;l;N,RL,j!QDRvqeCXONJhNfnU;DU-DUG?vECtAXvq$kqfNHHlI;EEAvW!T!LL?bgR'oI;,:gyvlojagUaMgyxv:-Ev,;uKHoeSD,; rBQNz$oIKHrVk!ny:KHAqNrD;wsoh3lUEtH?jdoe'3:Uj?o;wW;$UG$UEtUjpT!$c Q!ZqxiUGCB:zaEt$rllqegoCxNiNhfpDJLhZq3IS&hl!tC;U;EBB3rUL&TNr$rU;LqHwj;L?HTNjM?j?a;wTULCNS:ij;L&W;frct,DNDzT?joDlUqrmSFtWC$rrbDf?CfCegU;wTegnhWqrTCzL'$j;wNlN DKTUqrP$;'$AoD;w,axogCevBT;wjpZl?NP?zegH?B3Dr,llaULNj!bDNjXllegHCNjoAa;U;kgW;b ,!'llqv3IQPtCZtzT!thfTrblU;UfUUrNzjEthjUE;ZXNXkUYjl$U;vgjZbrF?lyNPC?Q;vqfR'$N:D'W;qxUrhN:DeNNXUakUEz3r&qi;C
------------------------------
step 100: train loss 2.0499, val loss 2.1563
------------------------------
ADY itlocthet tws le de fat be earrmideane ecer ne t, woho t's hin tofr aluyo'ssst of,
Tof t t yealllonarp g?


To ar thoo she lve ge grethes thofulllofl fifa r be Clle rdaree,
Whathelling gean ave huser e thofond f m yore tlinde trersthat, younoh?
py

Cr frsthe Rist t ishe latheavimbe ge je are hoffat t t, sthathee arelindld lli andlllll acl ttcle be fae spth.
Tobucaindse.
f Endalorst qun t ir e t alofan theadefledw avey arid inrdur twolll mup ard, clesert yofr, asttan fow oucoftillll t thin ss tofud yon t
------------------------------
step 200: train loss 1.6663, val loss 1.8703
------------------------------
herea yow hore ollde ge tthire of thavearn ide gllllowblmast achitw be s we tthite islle n lloure g putt offoor ye be ie he mettwstthe arve be idddathed o's, ung'dedlle be outt
The qut nlir,
Wathow youthow?
Herddlltint own ith?


Bie an on uthalllsthe thie s ats tvyoimime t rddighe oourdg tho's ave o p f qungooir quchut tt owe myoort ourdderse himithe t hoignve qut best oweeaton i I e asllllllff men youtsp Cig he it off or pst cowenchendinfllat igoulllllloow?

ckn ow ow?
F org t orst ame age owhthllick, ond
------------------------------
step 300: train loss 1.5378, val loss 1.7651
------------------------------
Thowlould ich o'dlldst yo be thendouth p.

Indeeembeeat.

Wit thiveree arst t me
Wandyoullmat h inddde impphmpy;
Ded, adddie
I'st ire y re t at he myo'Are heyost, memp, thouthemblerde
Be alltse ive me ff ust alman it utwat che I; cherenddddo wheremumse m,
t tho ade ale
Winddwhis
An o wallllit be be
Off y myore himph'st thide t thithouthe yere h,
Woursthem, theyout there avy
ghleereer?
In t imillddse ma it adde at are unt t a upo t cha wis yofffft an yoout yosthe pha eals anddde youthe aff ighe ist,
Wande it
------------------------------

While before using the rotary embeddings and just using learned position embeddings it looked like this:

step 0: train loss 5.7616, val loss 5.7801
------------------------------
;l;N,RL,j!QDRvqeCXONJhNfnU;DU-DUG?vECtAXvq$kqfNHHlI;EEAvW!T!LL?bgR'oI;,:gyvlojagUaMgyxv:-Ev,;uKHoeSD,; rBQNz$oIKHrVk!ny:KHAqNrD;wsoh3lUEtH?jdoe'3:Uj?o;wW;$UG$UEtUjpT!$c Q!ZqxiUGCB:zaEt$rllqegoCxNiNhfpDJLhZq3IS&hl!tC;U;EBB3rUL&TNr$rU;LqHwj;L?HTNjM?j?a;wTULCNS:ij;L&W;frct,DNDzT?joDlUqrmSFtWC$rrbDf?CfCegU;wTegnhWqrTCzL'$j;wNlN DKTUqrP$;'$AoD;w,axogCevBT;wjpZl?NP?zegH?B3Dr,llaULNj!bDNjXllegHCNjoAa;U;kgW;b ,!'llqv3IQPtCZtzT!thfTrblU;UfUUrNzjEthjUE;ZXNXkUYjl$U;vgjZbrF?lyNPC?Q;vqfR'$N:D'W;qxUrhN:DeNNXUakUEz3r&qi;C
------------------------------
step 100: train loss 2.4826, val loss 2.5239
------------------------------


T: youcthoveryo ime'd d t bealelour ff.


F t is, IStort w ildledes t SBythathethend stherde
Y:
Angreens e t, to avesethathe ly he,

Thumes wit wenkino en?

I be BRLI oncirel thaveerf, t edam ay'd wins d
Thoririf CO,-rshatind, frursthe upry whor fut bet estothe my!
MEOULI l thatrs.

BOLIOLI hef t herPr; Wisthathatharer ist hela

Alllerof ICETonel are,


LLULA:
I's FERKINI w CORIIUSThenor wast ber IDYO:

TIER:
NGHET:
Why IOR:
CIUSTI myenot hur:
MENIARDY:
WAy, s s tovef bey my s woouromy were ot.
May ayos s
------------------------------
step 200: train loss 2.2478, val loss 2.3003
------------------------------
TUSS:
site, fore the eantoress o athargher iche sises?
Thas dachows bonn wisthenon?
Anf he isureon:
Sursend whily s bedare foowiviss'd bly han hale pint ofe tomp?
Anou bed otivan'd sthenour falllou wourity s;
Thererthu nou knof dafe dak ondrep:
mowish urid t amyousup.

LULUCIO:
Ife, lleat tou shay so, be sullor, dop guresuant inted, dit,
Orvilimed, sede ho, wthepppe furenterist,
LA denatest miend by yours?

INGHENIUS:
O momet'd no:
Pe withist theere e myou the thod?

Nuranderm thentiof io.

COROROMIO:
If ch
------------------------------
step 300: train loss 2.0852, val loss 2.2030
------------------------------
QUETIS:
No ELlo VO:
Yo, ca--
Therestis re O, beare.

CORINTHUS:
To the athrer me theay's tr the here, ineroussws beld.

CLTENGO:
Clleay ow thave farim Eve heyour,
O, te sovetck, tho mabe, hert inte lis thing ang sino.
To the ee wim honne her urer toJubed
By stre foid bantce.'sen Bo nes eancence, ap hit
tand knecoble it the rener suptot: If so f councyou,
And iveionk, ur aus thor tist bestid co haly be besh,
My bucasediet t bespe the faute wh witw bo't grods dos,
Than lit bes rofur? the do ife cont.

COMPEST
------------------------------

You can see that the sampling is much better at a higher loss. But in my pytorch implementation of retnet and rotary embeddings, the sampling was not affected and trained much better.
P.S: In my pytorch implementation, the initial losses were around 4.3 which is expected as i have a vocab size of 65 and the initial loss is expected to be around ln(65) = 4.17, but in flax it is much larger. Although maybe not significant, what could be a reason for this?

[faresobeid]

** To clarify, in pytorch, the sampling and training drastically improved with rotary embeddings