In this work the invariant imbedding method has been used to develop an algorithm to study the torsional vibration of non-uniform systems. The algorithm is based on the propagation, reflection, and transmission of waves in a stepped waveguide and is part of a procedure to transform two-point boundary value problems in initial value problems. Based on this approach, a continuous model has been developed and a simple, versatile, and robust algorithm has been constructed to solve torsional vibration problems of non-uniform shafts with circular crosssections. The proposed solution algorithm was extensively evaluated through comparisons with analytical solutions and the finite element method. The results show that the proposed method can provide the exact solution for uniform shafts with concentrated elements and accurate results for a wide variety of torsional vibration problems. Systems with continuously varying geometry may be approximated by stepped shafts. The proposed method can also be used to study the dynamic behaviour of others stepped systems.